"""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 os
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)