"""The UCB module contains functions specific to 61A projects at UC Berkeley."""

import code
import functools
import inspect
import re
import signal
import sys
import timeit

def main(fn):
    """Call fn with command line arguments.  Used as a decorator.

    The main decorator marks the function that starts a program. For example,

    @main
    def my_run_function():
        # function body

    Use this instead of the typical __name__ == "__main__" predicate.
    """
    if inspect.stack()[1][0].f_locals['__name__'] == '__main__':
        args = sys.argv[1:] # Discard the script name from command line
        fn(*args) # Call the main function
    return fn

_PREFIX = ''
def trace(fn):
    """A decorator that prints a function's name, its arguments, and its return
    values each time the function is called. For example,

    @trace
    def compute_something(x, y):
        # function body
    """
    @functools.wraps(fn)
    def wrapped(*args, **kwds):
        global _PREFIX
        reprs = [repr(e) for e in args]
        reprs += [repr(k) + '=' + repr(v) for k, v in kwds.items()]
        log('{0}({1})'.format(fn.__name__, ', '.join(reprs)) + ':')
        _PREFIX += '    '
        try:
            result = fn(*args, **kwds)
            _PREFIX = _PREFIX[:-4]
        except Exception as e:
            log(fn.__name__ + ' exited via exception')
            _PREFIX = _PREFIX[:-4]
            raise
        # Here, print out the return value.
        log('{0}({1}) -> {2}'.format(fn.__name__, ', '.join(reprs), result))
        return result
    return wrapped


def log(message):
    """Print an indented message (used with trace)."""
    if type(message) is not str:
        message = str(message)
    print(_PREFIX + re.sub('\n', '\n' + _PREFIX, message))


def log_current_line():
    """Print information about the current line of code."""
    frame = inspect.stack()[1]
    log('Current line: File "{f[1]}", line {f[2]}, in {f[3]}'.format(f=frame))


def interact(msg=None):
    """Start an interactive interpreter session in the current environment.

    On Unix:
      <Control>-D exits the interactive session and returns to normal execution.
    In Windows:
      <Control>-Z <Enter> exists the interactive session and returns to normal
      execution.
    """
    # use exception trick to pick up the current frame
    try:
        raise None
    except:
        frame = sys.exc_info()[2].tb_frame.f_back

    # evaluate commands in current namespace
    namespace = frame.f_globals.copy()
    namespace.update(frame.f_locals)

    # exit on interrupt
    def handler(signum, frame):
        print()
        exit(0)
    signal.signal(signal.SIGINT, handler)

    if not msg:
        _, filename, line, _, _, _ = inspect.stack()[1]
        msg = 'Interacting at File "{0}", line {1} \n'.format(filename, line)
        msg += '    Unix:    <Control>-D continues the program; \n'
        msg += '    Windows: <Control>-Z <Enter> continues the program; \n'
        msg += '    exit() or <Control>-C exits the program'

    code.interact(msg, None, namespace)

def time_expr(expr, setup=None, imports=None, number=1000):
    """A convenience function for use with timeit.repeat.  Returns the
    minimum average per-iteration time of of 3 runs in which EXPR
    (a Python expression as a string) is executed NUMBER times.  Before
    executing each loop, executes SETUP (a Python expression as a string)
    and, if IMPORTS is present, executes an import of all these function names
    in IMPORTS (a list or string) from __main__."""
    if type(imports) is str:
        imports = "from __main__ import " + imports
    elif imports is not None:
        imports = "from __main__ import " + ", ".join(imports)
    if setup is None:
        setup = imports or ""
    elif imports is not None:
        setup = imports + "; " + setup
    return min(timeit.repeat(expr, setup, number=number)) / number

def time_expr_str(expr, setup=None, imports=None, number=1000):
    """A readable string representation of the result of
    time_expr(expr,setup,imports,number) in appropriate units."""
    t = time_expr(expr, setup=setup, imports=imports, number=number)
    if t < 0.001:
        return "{} usec".format(int(t * 1e6))
    elif t < 1.0:
        return "{} msec".format(int(t * 1000))
    else:
        return "{} sec".format(int(t))

def desc_time(expr, setup=None, imports=None, number=1000):
    """A description of the result and parameters of 
    time_expr(expr,setup,imports,number) in appropriate units."""

    t = time_expr_str(expr, setup=setup, imports=imports, number=number)
    return "{} loops, best of 3: {} per loop".format(number, t)