# Do not put a sh-bang line here (to allow different python builds to call the benchmarker).
import os
import sys
if sys.platform.startswith('win'):
    from time import clock as time
else:
    from time import time
from glob import glob
import psyco
import shelve
import gc


FASTREPS = 5
BASELINE = shelve.open('baseline.shl')

#2.5 Timing = NamedTuple('Timing', 'module func plaintime psycotime ratio')
# I don't have the energy to back-port this thing to 2.4  CT
class Timing(tuple):
    __repr__ = lambda self: '%-15s: %-30s plain: %7.0f/s     psyco: %10.0f/s     ratio: %6.2f' % self
    @property
    def module(self): return self[0]
    @property
    def func(self): return self[1]
    @property
    def plaintime(self): return self[2]
    @property
    def psycotime(self): return self[3]
    @property
    def ratio(self): return self[4]

DEBUG = False

def auto_time(f):
    """ run f until it has taken approx. 1/4 seconds, then
        estimate the number of repetitions needed for one second
    """
    _time = time
    k = 2
    while 1:
        start = _time()
        f(k)
        stop = _time()
        if stop - start >= 0.25:
            break
        k = 2*k
    used = stop - start
    res = int(k / used + .5)
    # goal is to have at least 5 runs or more
    assert res >= 5, ("%s's execution time is too long, "
                      "would be %d runs, only" % (f.__name__, res))
    return res

def odd(x):
    return x % 2 == 1
    
def median(lis):
    l = len(lis)
    h = l // 2
    if odd(l):
        return lis[h]
    else:
        return (lis[h] + lis[h+1]) / 2
        
def time_it(plain, compiled, once, unittest):
    _time = time
    result = []
    for func in (plain, compiled):
        loops = []
        run_each = once and 1 or 7
        runs = auto_time(func)
        if unittest:
            runs = 2
        for i in range(run_each):
            if DEBUG: print func,
            assert not __in_psyco__
            gc.collect()
            start = _time()
            func(runs)
            stop = _time()
            loops.append( runs / (stop - start) )
            if DEBUG: print "%0.5f" % loops[-1]
        if DEBUG: print ["  %0.5f" % x for x in loops]
        result.append(max(loops))
    plainloops, psycoloops = result        

    descr = plain.__doc__ or plain.__name__
    if descr.startswith('time_'):
        descr = descr[5:]
    
    return Timing((plain.__module__, descr, plainloops, psycoloops, psycoloops/plainloops))

def find_timing_modules():
    return sorted(glob('time_*.py'))

def time_modules(module_names, once=False, update=False, unittest=False):
    version = getattr(psyco, 'where_am_i', 'original psyco') + ' ' + getattr(psyco, 'when_was_i_born', '')
    print 'Running new timings with', version
    if update:
        BASELINE['version'] = version
    elif 'version' in BASELINE:
        print '            comparing to', BASELINE['version']
    print

    for module_name in module_names:
        if module_name.endswith('.py'):
            module_name = module_name[:-3]
        # the plain functions...
        mod = __import__(module_name)
        # ..and compiled versions.
        class compile_me(object):
            pass
        compile_me = compile_me() # cannot use a dict proxy
        # make sure everything on top-level gets compiled...
        psyco.full()
        execfile(module_name + '.py', compile_me.__dict__)
        # and also everything that will be called.
        psyco.bind(compile_me)
        psyco.stop()
        for varname in dir(mod):
            if not varname.startswith('time_'):
                continue
            plain = getattr(mod, varname)
            compiled = getattr(compile_me, varname)
            timing = time_it(plain, compiled, once, unittest)

            # print and either update or compare
            key = '%s:%s' % (timing.module, timing.func)
            if update:
                print timing
                BASELINE[key] = timing
            elif key in BASELINE:
                print 'new:', timing
                print 'old:', BASELINE[key]
                print
            else:
                print timing

if __name__ == '__main__':    
    from optparse import OptionParser

    parser = OptionParser()
    parser.add_option("-m", "--modules", action='store_true', dest="modules",
                      default=False, help="time listed modules")
    parser.add_option("-x", "--exclude", action='store_true', dest="exclude",
                      default=False, help="exclude from listed modules from search")
    parser.add_option("-q", "--quick", action='store_true', dest="once",
                      default=False, help="do each test only once, seven times"
                                          " quicker, but less accurate")
    parser.add_option("-t", "--unittest", action='store_true', dest="unittest",
                      default=False, help="run each bench only one time")
    parser.add_option("-u", "--update", action='store_true', dest="update",
                      default=False, help="update the base set of timings")
    parser.add_option("-v", "--verbose", action='store_true', dest="verbose",
                      default=False, help="print verbose timing output")
    opt, args = parser.parse_args()
    
    if opt.modules:
        modules = args
    elif opt.exclude:
        exclude = set(m for m in args if m.endswith('.py')) + set(
                      (m + '.py') for m in args if not m.endswith('.py') )
        modules = [m for m in find_timing_modules() if m not in exclude]
    else:
        modules = find_timing_modules()

    DEBUG = opt.verbose
    time_modules(modules, once=opt.once, update=opt.update,
                 unittest=opt.unittest)
    gc.collect()