1# -*- coding: utf-8 -*-
2# MolMod is a collection of molecular modelling tools for python.
3# Copyright (C) 2007 - 2019 Toon Verstraelen <Toon.Verstraelen@UGent.be>, Center
4# for Molecular Modeling (CMM), Ghent University, Ghent, Belgium; all rights
5# reserved unless otherwise stated.
6#
7# This file is part of MolMod.
8#
9# MolMod is free software; you can redistribute it and/or
10# modify it under the terms of the GNU General Public License
11# as published by the Free Software Foundation; either version 3
12# of the License, or (at your option) any later version.
13#
14# MolMod is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program; if not, see <http://www.gnu.org/licenses/>
21#
22# --
23
24
25from __future__ import print_function, division
26
27import sys
28import os
29import platform
30import datetime
31import getpass
32try:
33    from time import process_time
34except ImportError:
35    from time import clock as process_time
36import codecs
37import locale
38import functools
39from contextlib import contextmanager
40
41from molmod.units import kjmol, kcalmol, electronvolt, angstrom, nanometer, \
42    femtosecond, picosecond, amu, deg, gram, centimeter
43
44
45__all__ = ['ScreenLog', 'TimerGroup']
46
47
48class Unit(object):
49    def __init__(self, kind, conversion, notation, format):
50        self.kind = kind
51        self.conversion = conversion
52        self.notation = notation
53        self.format = format
54
55    def __call__(self, value):
56        return self.format % (value/self.conversion)
57
58
59class UnitSystem(object):
60    def __init__(self, *units):
61        self.units = units
62        # check for duplicates
63        for i0, unit0 in enumerate(self.units):
64            for unit1 in self.units[:i0]:
65                if unit0.kind == unit1.kind:
66                    raise ValueError('The unit of \'%s\' is encountered twice.' % unit0.kind)
67
68    def log_info(self, log):
69        if log.do_low:
70            with log.section('UNITS'):
71                log('The following units will be used below:')
72                log.hline()
73                log('Kind          Conversion               Format Notation')
74                log.hline()
75                for unit in self.units:
76                    log('%13s %21.15e %9s %s' % (unit.kind, unit.conversion, unit.format, unit.notation))
77                log.hline()
78                log('The internal data is divided by the corresponding conversion factor before it gets printed on screen.')
79
80    def apply(self, some):
81        for unit in self.units:
82            some.__dict__[unit.kind] = unit
83
84
85class ScreenLog(object):
86    # log levels
87    silent = 0
88    warning = 1
89    low = 2
90    medium = 3
91    high = 4
92    debug = 5
93
94    # screen parameters
95    margin = 8
96    width = 72
97
98    # unit systems
99    # TODO: the formats may need some tuning
100    joule = UnitSystem(
101        Unit('energy', kjmol, 'kJ/mol', '%10.1f'),
102        Unit('temperature', 1, 'K', '%10.1f'),
103        Unit('length', angstrom, 'A', '%10.4f'),
104        Unit('invlength', 1/angstrom, 'A^-1', '%10.5f'),
105        Unit('area', angstrom**2, 'A^2', '%10.3f'),
106        Unit('volume', angstrom**3, 'A^3', '%10.3f'),
107        Unit('time', femtosecond, 'fs', '%10.1f'),
108        Unit('mass', amu, 'amu', '%10.5f'),
109        Unit('charge', 1, 'e', '%10.5f'),
110        Unit('force', kjmol/angstrom, 'kJ/mol/A', '%10.1f'),
111        Unit('forceconst', kjmol/angstrom**2, 'kJ/mol/A**2', '%10.1f'),
112        Unit('velocity', angstrom/femtosecond, 'A/fs', '%10.5f'),
113        Unit('acceleration', angstrom/femtosecond**2, 'A/fs**2', '%10.5f'),
114        Unit('angle', deg, 'deg', '%10.5f'),
115        Unit('c6', 1, 'E_h*a_0**6', '%10.5f'),
116        Unit('c8', 1, 'E_h*a_0**6', '%10.5f'),
117        Unit('c12', 1, 'E_h*a_0**12', '%10.5f'),
118        Unit('diffconst', angstrom**2/picosecond, 'A**2/ps', '%10.5f'),
119        Unit('density', gram/centimeter**3, 'g/cm^3', '%10.3f'),
120    )
121    cal = UnitSystem(
122        Unit('energy', kcalmol, 'kcal/mol', '%10.2f'),
123        Unit('temperature', 1, 'K', '%10.1f'),
124        Unit('length', angstrom, 'A', '%10.4f'),
125        Unit('invlength', 1/angstrom, 'A^-1', '%10.5f'),
126        Unit('area', angstrom**2, 'A^2', '%10.3f'),
127        Unit('volume', angstrom**3, 'A^3', '%10.3f'),
128        Unit('time', femtosecond, 'fs', '%10.1f'),
129        Unit('mass', amu, 'amu', '%10.5f'),
130        Unit('charge', 1, 'e', '%10.5f'),
131        Unit('force', kcalmol/angstrom, 'kcal/mol/A', '%10.1f'),
132        Unit('forceconst', kcalmol/angstrom**2, 'kcal/mol/A**2', '%10.1f'),
133        Unit('velocity', angstrom/femtosecond, 'A/fs', '%10.5f'),
134        Unit('acceleration', angstrom/femtosecond**2, 'A/fs**2', '%10.5f'),
135        Unit('angle', deg, 'deg', '%10.5f'),
136        Unit('c6', 1, 'E_h*a_0**6', '%10.5f'),
137        Unit('c8', 1, 'E_h*a_0**6', '%10.5f'),
138        Unit('c12', 1, 'E_h*a_0**12', '%10.5f'),
139        Unit('diffconst', angstrom**2/femtosecond, 'A**2/fs', '%10.5f'),
140        Unit('density', gram/centimeter**3, 'g/cm^3', '%10.3f'),
141    )
142    solid = UnitSystem(
143        Unit('energy', electronvolt, 'eV', '%10.4f'),
144        Unit('temperature', 1, 'K', '%10.1f'),
145        Unit('length', angstrom, 'A', '%10.4f'),
146        Unit('invlength', 1/angstrom, 'A^-1', '%10.5f'),
147        Unit('area', angstrom**2, 'A^2', '%10.3f'),
148        Unit('volume', angstrom**3, 'A^3', '%10.3f'),
149        Unit('time', femtosecond, 'fs', '%10.1f'),
150        Unit('mass', amu, 'amu', '%10.5f'),
151        Unit('charge', 1, 'e', '%10.5f'),
152        Unit('force', electronvolt/angstrom, 'eV/A', '%10.1f'),
153        Unit('forceconst', electronvolt/angstrom**2, 'eV/A**2', '%10.1f'),
154        Unit('velocity', angstrom/femtosecond, 'A/fs', '%10.5f'),
155        Unit('acceleration', angstrom/femtosecond**2, 'A/fs**2', '%10.5f'),
156        Unit('angle', deg, 'deg', '%10.5f'),
157        Unit('c6', 1, 'E_h*a_0**6', '%10.5f'),
158        Unit('c8', 1, 'E_h*a_0**6', '%10.5f'),
159        Unit('c12', 1, 'E_h*a_0**12', '%10.5f'),
160        Unit('diffconst', angstrom**2/femtosecond, 'A**2/fs', '%10.5f'),
161        Unit('density', gram/centimeter**3, 'g/cm^3', '%10.3f'),
162    )
163    bio = UnitSystem(
164        Unit('energy', kcalmol, 'kcal/mol', '%10.2f'),
165        Unit('temperature', 1, 'K', '%10.1f'),
166        Unit('length', nanometer, 'nm', '%10.6f'),
167        Unit('area', nanometer**2, 'nm^2', '%10.4f'),
168        Unit('volume', nanometer**3, 'nanometer^3', '%10.1f'),
169        Unit('invlength', 1/nanometer, 'nm^-1', '%10.8f'),
170        Unit('time', picosecond, 'ps', '%10.4f'),
171        Unit('mass', amu, 'amu', '%10.5f'),
172        Unit('charge', 1, 'e', '%10.5f'),
173        Unit('force', kcalmol/angstrom, 'kcal/mol/A', '%10.5f'),
174        Unit('forceconst', kcalmol/angstrom**2, 'kcal/mol/A**2', '%10.5f'),
175        Unit('velocity', angstrom/picosecond, 'A/ps', '%10.5f'),
176        Unit('acceleration', angstrom/picosecond**2, 'A/ps**2', '%10.5f'),
177        Unit('angle', deg, 'deg', '%10.5f'),
178        Unit('c6', 1, 'E_h*a_0**6', '%10.5f'),
179        Unit('c8', 1, 'E_h*a_0**6', '%10.5f'),
180        Unit('c12', 1, 'E_h*a_0**12', '%10.5f'),
181        Unit('diffconst', nanometer**2/picosecond, 'nm**2/ps', '%10.2f'),
182        Unit('density', gram/centimeter**3, 'g/cm^3', '%10.3f'),
183    )
184    atomic = UnitSystem(
185        Unit('energy', 1, 'E_h', '%10.6f'),
186        Unit('temperature', 1, 'K', '%10.1f'),
187        Unit('length', 1, 'a_0', '%10.5f'),
188        Unit('invlength', 1, 'a_0^-1', '%10.5f'),
189        Unit('area', 1, 'a_0^2', '%10.3f'),
190        Unit('volume', 1, 'a_0^3', '%10.3f'),
191        Unit('time', 1, 'aut', '%10.1f'),
192        Unit('mass', 1, 'aum', '%10.1f'),
193        Unit('charge', 1, 'e', '%10.5f'),
194        Unit('force', 1, 'E_h/a_0', '%10.5f'),
195        Unit('forceconst', 1, 'E_h/a_0**2', '%10.5f'),
196        Unit('velocity', 1, 'a_0/aut', '%10.5f'),
197        Unit('acceleration', 1, 'a_0/aut**2', '%10.5f'),
198        Unit('angle', 1, 'rad', '%10.7f'),
199        Unit('c6', 1, 'E_h*a_0**6', '%10.5f'),
200        Unit('c8', 1, 'E_h*a_0**6', '%10.5f'),
201        Unit('c12', 1, 'E_h*a_0**12', '%10.5f'),
202        Unit('diffconst', 1, 'a_0**2/aut', '%10.2f'),
203        Unit('density', 1, 'aum/a_0^3', '%10.3f'),
204    )
205
206
207    def __init__(self, name, version, head_banner, foot_banner, timer, f=None):
208        self.name = name
209        self.version = version
210        self.head_banner = head_banner
211        self.foot_banner = foot_banner
212        self.timer = timer
213
214        self._active = False
215        self._level = self.medium
216        self.unitsys = self.joule
217        self.unitsys.apply(self)
218        self.prefix = ' '*(self.margin-1)
219        self._last_used_prefix = None
220        self.stack = []
221        self.add_newline = False
222        if f is None:
223            _file = sys.stdout
224        else:
225            _file = f
226        self.set_file(_file)
227
228    do_warning = property(lambda self: self._level >= self.warning)
229    do_low = property(lambda self: self._level >= self.low)
230    do_medium = property(lambda self: self._level >= self.medium)
231    do_high = property(lambda self: self._level >= self.high)
232    do_debug = property(lambda self: self._level >= self.debug)
233
234    def _pass_color_code(self, code):
235        if self._file.isatty():
236            return code
237        else:
238            return ''
239
240    reset =     property(functools.partial(_pass_color_code, code="\033[0m"))
241    bold =      property(functools.partial(_pass_color_code, code="\033[01m"))
242    teal =      property(functools.partial(_pass_color_code, code="\033[36;06m"))
243    turquoise = property(functools.partial(_pass_color_code, code="\033[36;01m"))
244    fuchsia =   property(functools.partial(_pass_color_code, code="\033[35;01m"))
245    purple =    property(functools.partial(_pass_color_code, code="\033[35;06m"))
246    blue =      property(functools.partial(_pass_color_code, code="\033[34;01m"))
247    darkblue =  property(functools.partial(_pass_color_code, code="\033[34;06m"))
248    green =     property(functools.partial(_pass_color_code, code="\033[32;01m"))
249    darkgreen = property(functools.partial(_pass_color_code, code="\033[32;06m"))
250    yellow =    property(functools.partial(_pass_color_code, code="\033[33;01m"))
251    brown =     property(functools.partial(_pass_color_code, code="\033[33;06m"))
252    red =       property(functools.partial(_pass_color_code, code="\033[31;01m"))
253
254    def set_file(self, f):
255        # Wrap sys.stdout into a StreamWriter to allow writing unicode.
256        self._file = f
257        self.add_newline = False
258
259    def set_level(self, level):
260        if level < self.silent or level > self.debug:
261            raise ValueError('The level must be one of the ScreenLog attributes.')
262        self._level = level
263
264    def __call__(self, *words):
265        s = u' '.join(str(w) for w in words)
266        if not self.do_warning:
267            raise RuntimeError('The runlevel should be at least warning when logging.')
268        if not self._active:
269            prefix = self.prefix
270            self.print_header()
271            self.prefix = prefix
272        if self.add_newline and self.prefix != self._last_used_prefix:
273            self._file.write(u'\n')
274            self.add_newline = False
275        # Check for alignment code '&'
276        pos = s.find(u'&')
277        if pos == -1:
278            lead = u''
279            rest = s
280        else:
281            lead = s[:pos] + ' '
282            rest = s[pos+1:]
283        width = self.width - len(lead)
284        if width < self.width//2:
285            raise ValueError('The lead may not exceed half the width of the terminal.')
286        # break and print the line
287        first = True
288        while len(rest) > 0:
289            if len(rest) > width:
290                pos = rest.rfind(' ', 0, width)
291                if pos == -1:
292                    current = rest[:width]
293                    rest = rest[width:]
294                else:
295                    current = rest[:pos]
296                    rest = rest[pos:].lstrip()
297            else:
298                current = rest
299                rest = u''
300            self._file.write(u'%s %s%s\n' % (self.prefix, lead, current))
301            if first:
302                lead = u' '*len(lead)
303                first = False
304        self._last_used_prefix = self.prefix
305
306    def warn(self, *words):
307        self(u'WARNING!!&'+u' '.join(str(w) for w in words))
308
309    def hline(self, char='~'):
310        self(char*self.width)
311
312    def center(self, *words, **kwargs):
313        if len(kwargs) == 0:
314            edge = ''
315        elif len(kwargs) == 1:
316            if 'edge' not in kwargs:
317                raise TypeError('Only one keyword argument is allowed, that is edge')
318            edge = kwargs['edge']
319        else:
320            raise TypeError('Too many keyword arguments. Should be at most one.')
321        s = u' '.join(str(w) for w in words)
322        if len(s) + 2*len(edge) > self.width:
323            raise ValueError('Line too long. center method does not support wrapping.')
324        self('%s%s%s' % (edge, s.center(self.width-2*len(edge)), edge))
325
326    def blank(self):
327        self._file.write(u'\n')
328
329    def _enter(self, prefix):
330        if len(prefix) > self.margin-1:
331            raise ValueError('The prefix must be at most %s characters wide.' % (self.margin-1))
332        self.stack.append(self.prefix)
333        self.prefix = prefix.upper().rjust(self.margin-1, ' ')
334        self.add_newline = True
335
336    def _exit(self):
337        self.prefix = self.stack.pop(-1)
338        if self._active:
339            self.add_newline = True
340
341    @contextmanager
342    def section(self, prefix):
343        self._enter(prefix)
344        try:
345            yield
346        finally:
347            self._exit()
348
349    def set_unitsys(self, unitsys):
350        self.unitsys = unitsys
351        self.unitsys.apply(self)
352        if self._active:
353            self.unitsys.log_info()
354
355    def print_header(self):
356        # Suppress any logging as soon as an exception is not caught.
357        def excepthook_wrapper(type, value, traceback):
358            self.set_level(self.silent)
359            sys.__excepthook__(type, value, traceback)
360        sys.excepthook = excepthook_wrapper
361
362        if self.do_warning and not self._active:
363            self._active = True
364            print(self.head_banner, file=self._file)
365            self._print_basic_info()
366            self.unitsys.log_info(self)
367
368    def print_footer(self):
369        if self.do_warning and self._active:
370            self._print_basic_info()
371            self.timer._stop('Total')
372            self.timer.report(self)
373            print(self.foot_banner, file=self._file)
374
375    def _print_basic_info(self):
376        if self.do_low:
377            with self.section('ENV'):
378                self('User:          &' + getpass.getuser())
379                self('Platform:      &' + platform.platform())
380                self('Time:          &' + datetime.datetime.now().isoformat())
381                self('Python version:&' + sys.version.replace('\n', ''))
382                self('%s&%s' % (('%s version:' % self.name).ljust(15), self.version))
383                self('Current Dir:   &' + os.getcwd())
384                self('Command line:  &' + ' '.join(sys.argv))
385
386
387class Timer(object):
388    def __init__(self):
389        self.cpu = 0.0
390        self._start = None
391
392    def start(self):
393        assert self._start is None
394        self._start = process_time()
395
396    def stop(self):
397        assert self._start is not None
398        self.cpu += process_time() - self._start
399        self._start = None
400
401
402class SubTimer(object):
403    def __init__(self, label):
404        self.label = label
405        self.total = Timer()
406        self.own = Timer()
407
408    def start(self):
409        self.total.start()
410        self.own.start()
411
412    def start_sub(self):
413        self.own.stop()
414
415    def stop_sub(self):
416        self.own.start()
417
418    def stop(self):
419        self.own.stop()
420        self.total.stop()
421
422
423class TimerGroup(object):
424    def __init__(self):
425        self.parts = {}
426        self._stack = []
427        self._start('Total')
428
429    def reset(self):
430        for timer in self.parts.values():
431            timer.total.cpu = 0.0
432            timer.own.cpu = 0.0
433
434    @contextmanager
435    def section(self, label):
436        self._start(label)
437        try:
438            yield
439        finally:
440            self._stop(label)
441
442    def _start(self, label):
443        # get the right timer object
444        timer = self.parts.get(label)
445        if timer is None:
446            timer = SubTimer(label)
447            self.parts[label] = timer
448        # start timing
449        timer.start()
450        if len(self._stack) > 0:
451            self._stack[-1].start_sub()
452        # put it on the stack
453        self._stack.append(timer)
454
455    def _stop(self, label):
456        timer = self._stack.pop(-1)
457        assert timer.label == label
458        timer.stop()
459        if len(self._stack) > 0:
460            self._stack[-1].stop_sub()
461
462    def get_max_own_cpu(self):
463        result = None
464        for part in self.parts.values():
465            if result is None or result < part.own.cpu:
466                result = part.own.cpu
467        return result
468
469    def report(self, log):
470        max_own_cpu = self.get_max_own_cpu()
471        #if max_own_cpu == 0.0:
472        #    return
473        with log.section('TIMER'):
474            log('Overview of CPU time usage.')
475            log.hline()
476            log('Label             Total      Own')
477            log.hline()
478            bar_width = log.width-33
479            for label, timer in sorted(self.parts.items()):
480                #if timer.total.cpu == 0.0:
481                #    continue
482                if max_own_cpu > 0:
483                    cpu_bar = "W"*int(timer.own.cpu/max_own_cpu*bar_width)
484                else:
485                    cpu_bar = ""
486                log('%14s %8.1f %8.1f %s' % (
487                    label.ljust(14),
488                    timer.total.cpu, timer.own.cpu, cpu_bar.ljust(bar_width),
489                ))
490            log.hline()
491