import numpy as np
import matplotlib.pyplot as plt
import timeit
import json
import rapidjson
import pickle


def mixed_dict(z):
    for x in range(z):
        d['tests' + str(x)] = 'this is a test ' + str(x)
        d['testi' + str(x)] = int(x)
        d['testf' + str(x)] = float(x)
        d[str(x)] = True


def arrays(z):
    d['data'] = []
    d['data2'] = []
    for x in range(z):
        d['data'].append(x)
        d['data2'].append(x)


for gen in (mixed_dict, arrays):

    d = {}

    means_json = []
    means_rapidjson = []
    means_pickle = []

    def test_json():
        json.dumps(d)

    def test_rapidjson():
        rapidjson.dumps(d)

    def test_pickle():
        pickle.dumps(d)

    tests = [test_json, test_rapidjson, test_pickle]

    n = 100

    for t in (10, 100, 1000, 10000):
        d.clear()
        gen(t)
        means_json.append(1 / timeit.timeit(stmt=test_json, number=n) * t)
        means_rapidjson.append(1 /
                               timeit.timeit(stmt=test_rapidjson, number=n) * t)
        means_pickle.append(1 / timeit.timeit(stmt=test_pickle, number=n) * t)

    # data to plot
    n_groups = 4

    # create plot
    fig, ax = plt.subplots()
    index = np.arange(n_groups)
    bar_width = 0.2
    opacity = 0.8

    rects1 = plt.bar(index,
                     means_json,
                     bar_width,
                     alpha=opacity,
                     color='b',
                     label='JSON')

    rects2 = plt.bar(index + bar_width,
                     means_rapidjson,
                     bar_width,
                     alpha=opacity,
                     color='g',
                     label='RapidJSON')

    rects3 = plt.bar(index + bar_width * 2,
                     means_pickle,
                     bar_width,
                     alpha=opacity,
                     color='orange',
                     label='Pickle')

    plt.xlabel('Data size')
    plt.ylabel('Score (serializations/sec * data_size)')
    plt.title('Score by algo (higher is better), data: ' + gen.__name__)
    plt.xticks(index + bar_width, ('10', '100', '1000', '10000'))
    plt.legend()

    plt.tight_layout()
    plt.show()