'''
  Module for usefull decorators e.g. for logging function calls, input and output.
'''

import sys
import inspect
import logging
import threading
from time import time
from io import StringIO

#
# Help functions adopted from Michele Simionato's decorator module 
# http://www.phyast.pitt.edu/~micheles/python/decorator.zip
#

def getinfo(func):
    """
      Returns an info dictionary containing:
      - name (the name of the function : str)
      - argnames (the names of the arguments : list)
      - defaults (the values of the default arguments : tuple)
      - signature (the signature : str)
      - doc (the docstring : str)
      - module (the module name : str)
      - dict (the function __dict__ : str)
    """
    assert inspect.ismethod(func) or inspect.isfunction(func)
    signature=inspect.signature(func)
    output=dict(name=func.__name__, argnames=signature.parameters.keys(), signature=str(signature),
                defaults=func.__defaults__, doc=func.__doc__,
                module=func.__module__, dict=func.__dict__,
                globals=func.__globals__, closure=func.__closure__)
    return output

def update_wrapper(wrapper, wrapped, create=False):
    """
      An improvement over functools.update_wrapper. By default it works the
      same, but if the 'create' flag is set, generates a copy of the wrapper
      with the right signature and update the copy, not the original.
      Moreovoer, 'wrapped' can be a dictionary with keys 'name', 'doc', 'module',
      'dict', 'defaults'.
    """
    if isinstance(wrapped, dict):
        infodict=wrapped
    else:  # assume wrapped is a function
        infodict=getinfo(wrapped)
    assert not '_wrapper_' in infodict["argnames"], \
        '"_wrapper_" is a reserved argument name!'
    if create:  # create a brand new wrapper with the right signature
        src="lambda %(signature)s: _wrapper_(%(signature)s)"%infodict
        # import sys; print >> sys.stderr, src # for debugging purposes
        wrapper=eval(src, dict(_wrapper_=wrapper))
    try:
        wrapper.__name__=infodict['name']
    except:  # Python version < 2.4
        pass
    wrapper.__doc__=infodict['doc']
    wrapper.__module__=infodict['module']
    wrapper.__dict__.update(infodict['dict'])
    wrapper.__defaults__=infodict['defaults']
    return wrapper

# the real meat is here
def _decorator(caller, func):
    infodict=getinfo(func)
    argnames=infodict['argnames']
    assert not ('_call_' in argnames or '_func_' in argnames), \
        'You cannot use _call_ or _func_ as argument names!'
    src="lambda %(signature)s: _call_(_func_, %(signature)s)"%infodict
    dec_func=eval(src, dict(_func_=func, _call_=caller))
    return update_wrapper(dec_func, func)

def decorator(caller):
    """
    General purpose decorator factory: takes a caller function as
    input and returns a decorator with the same attributes.
    """
    return update_wrapper(lambda f: _decorator(caller, f), caller)

########################### Decorators for logging ##################################

def _nicetype(item):
    if 'array' in type(item).__name__:
        return type(item).__name__+("[%s]"%str(item.shape))
    elif type(item) in [list, tuple]:
        return type(item).__name__+('[%i]'%len(item))
    else:
        return type(item).__name__

class DecoratorLogger(logging.getLoggerClass()):
    '''
      A logger that makes sure the actual function definition filename, lineno and function name
      is used for logging.
    '''

    if sys.version_info[0:2]>=(3, 2):  # sinfo was introduced in python 3.2
        def makeRecord(self, name, lvl, fn, lno, msg, args, exc_info, func=None, extra=None, sinfo=None):
            if extra is None:
                return logging.getLoggerClass().makeRecord(self, name, lvl, fn, lno,
                                                           msg, args, exc_info, func=func, extra=None, sinfo=sinfo)
            else:
                return logging.getLoggerClass().makeRecord(self, name, lvl, extra['name'], extra['lno'],
                                                           msg, args, exc_info, func=extra['func'], extra=None,
                                                           sinfo=sinfo)
    else:
        def makeRecord(self, name, lvl, fn, lno, msg, args, exc_info, func=None, extra=None):
            if extra is None:
                return logging.getLoggerClass().makeRecord(self, name, lvl, fn, lno,
                                                           msg, args, exc_info, func=func, extra=None)
            else:
                return logging.getLoggerClass().makeRecord(self, name, lvl, extra['name'], extra['lno'],
                                                           msg, args, exc_info, func=extra['func'], extra=None)

# only use the debug decorators in debugging runs.
old_class=logging.getLoggerClass()
logging.setLoggerClass(DecoratorLogger)
logger=logging.getLogger('deco')
logger.setLevel(logging.DEBUG)
null_handler=logging.StreamHandler(StringIO())
null_handler.setLevel(logging.CRITICAL)
logger.addHandler(null_handler)
logging.setLoggerClass(old_class)

def _logformat(msg, decname, func):
    if hasattr(func, 'im_func'):
        co=func.__func__.__code__
    else:
        co=func.__code__
    fname=co.co_filename
    lno=co.co_firstlineno
    logger.debug(msg, extra={'func': '@'+decname, 'name': fname, 'lno': lno})

@decorator
def log_call(func, *args, **kw):
    '''
      Decorator to log just the method call.
    '''
    if logging.root.getEffectiveLevel()>logging.DEBUG: return func(*args, **kw)
    infodict=getinfo(func)
    if len(infodict['argnames'])>0 and infodict['argnames'][0]=='self':
        _logformat('%s.%s.%s <%s>'%(infodict['module'],
                                    args[0].__class__.__name__, infodict['name'], repr(threading.current_thread())),
                   'log_call', func)
    else:
        _logformat('%s.%s <%s>'%(infodict['module'], infodict['name'], repr(threading.current_thread())), 'log_call',
                   func)
    return func(*args, **kw)

@decorator
def log_input(func, *args, **kw):
    '''
      Decorator to log a method call with input.
    '''
    if logging.root.getEffectiveLevel()>logging.DEBUG: return func(*args, **kw)
    infodict=getinfo(func)
    if hasattr(func, 'im_func'):
        logstr=' call %s.%s('%(args[0].__class__.__name__, infodict['name'])
        method=True
        info_len=len('%s.%s('%(args[0].__class__.__name__, infodict['name']))
    else:
        logstr='call %s('%(infodict['name'])
        method=False
        info_len=len('%s('%(infodict['name']))

    namelen=max([len(name)+3 for name in infodict['argnames']]+[len(name) for name in list(kw.keys())])
    for i, arg in enumerate(args):
        if i==0 and (method or infodict['name']=='__init__'):
            continue
        value=repr(arg)
        value_split=value.splitlines()
        if len(value_split)>5:
            value="\n".join(value_split[:2]+[' '*int(len(value_split[1])/2)+'...']+value_split[-2:])
        value.replace('\n', '\n'+' '*int(namelen+7))
        logstr+='\n'+('%% %is= %%s (%%s)'%(namelen+3))%(infodict['argnames'][i], value, _nicetype(arg))

    for key, arg in sorted(kw.items()):
        value=repr(arg)
        value_split=value.splitlines()
        if len(value_split)>5:
            value="\n".join(value_split[:2]+[' '*int(len(value_split[1])/2)+'...']+value_split[-2:])
        value.replace('\n', '\n'+' '*int(namelen+4))
        logstr+='\n'+('KW %% %ss= %%s (%%s)'%namelen)%(key, repr(value), _nicetype(arg))
    logstr+='\n)'
    logstr=logstr.replace('\n', '\n'+' '*int(info_len+10))
    logstr=logstr.replace('call ', 'call \n'+' '*10)

    _logformat(logstr, 'log_input', func)
    return func(*args, **kw)

@decorator
def log_output(func, *args, **kw):
    '''
      Decorator to log a method call with output. If combined with log_input
      the input is logged at the time before the call and the output after.
    '''
    if logging.root.getEffectiveLevel()>logging.DEBUG: return func(*args, **kw)
    output=func(*args, **kw)
    infodict=getinfo(func)
    if len(infodict['argnames'])>0 and infodict['argnames'][0]=='self':
        logstr='return from %s.%s'%(args[0].__class__.__name__, infodict['name'])
    else:
        logstr='return from %s'%(infodict['name'])
    value=repr(output)
    value_split=value.splitlines()
    if len(value_split)>5:
        value="\n".join(value_split[:2]+[' '*int(len(value_split[1])/2)+'...']+value_split[-2:])
    logstr+='\n-> %15s (%s)'%(value, _nicetype(output))
    logstr=logstr.replace('\n', '\n'+' '*44)
    _logformat(logstr, 'log_output', func)
    return output

@decorator
def log_both(func, *args, **kw):
    '''
      Decoratore to log a method call with input and output.
    '''
    if logging.root.getEffectiveLevel()>logging.DEBUG: return func(*args, **kw)
    infodict=getinfo(func)
    if hasattr(func, 'im_func'):
        logstr=' call %s.%s('%(args[0].__class__.__name__, infodict['name'])
        method=True
        info_len=len('%s.%s('%(args[0].__class__.__name__, infodict['name']))
    else:
        logstr='call %s('%(infodict['name'])
        method=False
        info_len=len('%s('%(infodict['name']))

    namelen=max([len(name)+3 for name in infodict['argnames']]+[len(name) for name in list(kw.keys())])
    for i, arg in enumerate(args):
        if i==0 and (method or infodict['name']=='__init__'):
            continue
        value=repr(arg)
        value_split=value.splitlines()
        if len(value_split)>5:
            value="\n".join(value_split[:2]+[' '*int(len(value_split[1])/2)+'...']+value_split[-2:])
        value.replace('\n', '\n'+' '*int(namelen+7))
        logstr+='\n'+('%% %is= %%s (%%s)'%(namelen+3))%(infodict['argnames'][i], value, _nicetype(arg))

    for key, arg in sorted(kw.items()):
        value=repr(arg)
        value_split=value.splitlines()
        if len(value_split)>5:
            value="\n".join(value_split[:2]+[' '*int(len(value_split[1])/2)+'...']+value_split[-2:])
        value.replace('\n', '\n'+' '*int(namelen+4))
        logstr+='\n'+('KW %% %ss= %%s (%%s)'%namelen)%(key, repr(value), _nicetype(arg))
    logstr+='\n)'
    logstr=logstr.replace('\n', '\n'+' '*int(info_len+10))
    logstr=logstr.replace('call ', 'call \n'+' '*10)

    _logformat(logstr, 'log_input', func)
    # call the function
    output=func(*args, **kw)
    if len(infodict['argnames'])>0 and infodict['argnames'][0]=='self':
        logstr='return from %s.%s'%(args[0].__class__.__name__, infodict['name'])
    else:
        logstr='return from %s'%(infodict['name'])
    value=repr(output)
    value_split=value.splitlines()
    if len(value_split)>5:
        value="\n".join(value_split[:2]+[' '*int(len(value_split[1])/2)+'...']+value_split[-2:])
    logstr+='\n-> %s (%s)'%(value, _nicetype(output))
    logstr=logstr.replace('\n', '\n'+' '*int(info_len+10))
    _logformat(logstr, 'log_output', func)
    return output

timings={}

@decorator
def time_call(func, *args, **kw):
    '''
      Decorator to log just the method call.
    '''
    if logging.root.getEffectiveLevel()>logging.DEBUG: return func(*args, **kw)
    name=func.__name__
    start=time()
    output=func(*args, **kw)
    runtime=time()-start
    if not name in timings:
        timings[name]=(0., 0.)
    avg, calls=timings[name]
    timings[name]=((avg*calls+runtime)/(calls+1), calls+1)
    return output

########################## General decorators ###############################

class check_input(object):
    '''
      Decorator checking the input to a function.

      Usage:
      @check_input([type1, type2, ...], try_convert=True)
      def your_function(var1, var2, ...):
          ... function text ...
    '''
    try_convert=True
    types=[]

    def __init__(self, types, try_convert=True):
        self.types=types
        self.try_convert=try_convert

    def __call__(self, func):
        infodict=getinfo(func)
        argnames=infodict['argnames']
        assert not ('_call_' in argnames or '_func_' in argnames), \
            'You cannot use _call_ or _func_ as argument names!'
        src="def dec_func(%(signature)s):"%infodict
        for i, typei in enumerate(self.types):
            src+='\n  if type(%s).__name__!="%s":'%(argnames[i], typei.__name__)
            if self.try_convert:
                src+='\n    try:'
                src+='\n      %s=%s(%s)'%(argnames[i], typei.__name__, argnames[i])
                src+='\n    except:'
                src+='\n      raise ValueError("type of %s is not %s")'%(argnames[i], typei.__name__)
            else:
                src+='\n    raise ValueError("type of %s is not %s")'%(argnames[i], typei.__name__)
        src+='\n  return _func_(%(signature)s)'%infodict
        exec_dict=dict(_func_=func, _call_=self.__call__)
        exec(src, exec_dict)
        return update_wrapper(exec_dict['dec_func'], func)
