"""This module contains code from
Think Python by Allen B. Downey
http://thinkpython.com

Copyright 2012 Allen B. Downey
License: GNU GPLv3 http://www.gnu.org/licenses/gpl.html

This is 3.6 exercise 6:

"Test the performance of LinearMap, BetterMap and HashMap; see if you can
characterize their order of growth.  You can download my map implementations
from thinkcomplex.com/Map.py, and the code I used in this section from
thinkcomplex.com/listsum.py.

You will have to find a range of n that is big enough to show asymptotic
behavior, but small enough to run quickly."

"""
import random
import timeit

import matplotlib.pyplot as pyplot

from ch3ex4 import LinearMap, BetterMap, HashMap


MAP = None
VALUES = None

def test_map():

    MAP.get(random.choice(VALUES))


def test_func():

    elapsed = timeit.timeit('test_map()',
        setup='from __main__ import test_map', number=100)
    return elapsed


def test(map_class):
    """Tests the given function with a range of values for n.

    Returns a list of ns and a list of run times.
    """
    global MAP, VALUES

    klass = eval(map_class)

    d = {'LinearMap': 100000, 'BetterMap': 100000, 'HashMap': 100000}
    factor = d[map_class]

    # test (f) over a range of values for (n)
    ns = []
    ts = []
    for i in range(2, 25):
        n = factor * i

        MAP = klass()
        VALUES = range(n)
        random.shuffle(VALUES)
        for v in VALUES:
            MAP.add(v, v)

        t = test_func()

        print n, t
        ns.append(n)
        ts.append(t)

    return ns, ts


def plot(ns, ts, label, color='blue', exp=1.0):
    """Plots data and a fitted curve.

    ns: sequence of n (problem size)
    ts: sequence of t (run time)
    label: string label for the data curve
    color: string color for the data curve
    exp: exponent (slope) for the fitted curve
    """
    tfit = fit(ns, ts, exp)
    pyplot.plot(ns, tfit, color='0.7', linewidth=2, linestyle='dashed')
    pyplot.plot(ns, ts, label=label, color=color, linewidth=3)


def fit(ns, ts, exp=1.0, index=-1):
    """Fits a curve with the given exponent.

    Use the given index as a reference point, and scale all other
    points accordingly.
    """
    nref = ns[index]
    tref = ts[index]

    tfit = []
    for n in ns:
        ratio = float(n) / nref
        t = ratio**exp * tref
        tfit.append(t)

    return tfit


def save(root, exts=['eps', 'pdf']):
    """Saves the current figure in the given formats.

    root: string filename root
    exts: list of string extensions (specifying output formats).
    """
    for ext in exts:
        filename = '%s.%s' % (root, ext)
        print 'Writing', filename
        pyplot.savefig(filename)


def make_fig(funcs, scale='log', exp=1.0, filename=None):
    pyplot.clf()
    pyplot.xscale(scale)
    pyplot.yscale(scale)
    pyplot.title('')
    pyplot.xlabel('n')
    pyplot.ylabel('run time (s)')

    colors = ['blue', 'orange', 'green']
    for func, color in zip(funcs, colors):
        data = test(func)
        plot(*data, label=func, color=color, exp=exp)

    pyplot.legend(loc=4)

    if filename:
        save(filename)
    else:
        pyplot.show()


def main(script):
    make_fig(['LinearMap', 'BetterMap', 'HashMap'], scale='log', exp=1.0)


if __name__ == '__main__':
    import sys
    main(*sys.argv)