#! /usr/bin/env python
# Expects Python 3, numpy, and matplotlib


import datetime

import numpy as np
import matplotlib
matplotlib.use('SVG')
import matplotlib.pyplot as plt

MAX = 2 * 10**7
STEP = 2 * 10**4


def main():
    test_subject = set()
    insert_results = timed_calls(insert_step(test_subject), MAX, STEP)
    remove_results = timed_calls(remove_step(test_subject), MAX, STEP)
    assert len(test_subject) == 0, "Inserts should match removes"

    save_plot(
            [count for count, duration in insert_results],
            [duration for count, duration in insert_results],
            "Size", "Duration (s)",
            "Time to Insert a Batch of Items, vs Size")

    save_plot(
            [count for count, duration in insert_results],
            [duration for count, duration in amortize(insert_results)],
            "Size", "Duration (s)",
            "Amortized Time to Insert an Item, vs Size")

    save_plot(
            [count for count, duration in remove_results],
            [duration for count, duration in remove_results],
            "Total Removed", "Duration (s)",
            "Time to Remove a Batch of Items vs Total Removed")

    save_plot(
            [count for count, duration in remove_results],
            [duration for count, duration in amortize(remove_results)],
            "Total Removed", "Duration (s)",
            "Amortized Time to Remove an Item vs Total Removed")

def timed_calls(callback, target, stepsize):
    results = []
    for i in range(0, target, stepsize):
        before = datetime.datetime.now()
        callback(i, i + stepsize)
        after = datetime.datetime.now()
        results.append((i + stepsize, (after - before).total_seconds()))
    return results

def insert_step(container):
    def insert(start, stop):
        for i in range(start, stop):
            container.add(str(i))
    return insert

def remove_step(container):
    def remove(start, stop):
        for i in range(start, stop):
            container.remove(str(i))
    return remove

def amortize(records):
    results = []
    cumulative_total = 0
    for count, duration in records:
        cumulative_total += duration
        results.append((count, cumulative_total / count))
    return results

def save_plot(x, y, xlabel, ylabel, title):
    # Points
    plt.scatter(x, y)
    line, residual, _, _, _ = np.polyfit(x, y, 1, full=True)
    plt.plot(x, np.poly1d(line)(x),
            label="slope=%.2e,RMSE=%.2e" % (line[0], residual[0]/len(x)))
    plt.legend()
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.title(title)
    plt.savefig(title)
    plt.clf()
    with open(title + ".tab.txt", 'w') as tablefile:
        for point in zip(x, y):
            tablefile.write("%.2e\t%.2e\n" % point)

if __name__ == '__main__':
    main()