#!/usr/bin/env python3
# SPDX-License-Identifier: (Apache-2.0 OR MIT)

import collections
import io
import json
import os
import sys

from tabulate import tabulate
import matplotlib.pyplot as plt

LIBRARIES = ('orjson', 'ujson', 'rapidjson', 'json')

COLOR = ('blue', 'green', 'red', 'blue')

def aggregate():
    benchmarks_dir = os.path.join('.benchmarks', os.listdir('.benchmarks')[0])
    res = collections.defaultdict(dict)
    for filename in os.listdir(benchmarks_dir):
        with open(os.path.join(benchmarks_dir, filename), 'r') as fileh:
            data = json.loads(fileh.read())

        for each in data['benchmarks']:
            res[each['group']][each['extra_info']['lib']] = {
                'data': [
                    val * 1000 for val in each['stats']['data']
                ],
                'median': each['stats']['median'] * 1000,
                'ops': each['stats']['ops'],
            }
    return res

def box(obj):
    for group, val in sorted(obj.items()):
        data = []
        for lib in LIBRARIES:
            data.append(val[lib]['data'])
        fig = plt.figure(1, figsize=(9, 6))
        ax = fig.add_subplot(111)
        bp = ax.boxplot(data, vert=False, labels=LIBRARIES)
        ax.set_xlim(left=0)
        ax.set_xlabel('milliseconds')
        plt.title(group)
        plt.savefig('doc/{}.png'.format(group.replace(' ', '_').replace('.json', '')))
        plt.close()

def tab(obj):
    buf = io.StringIO()
    headers = ('Library', 'Median (milliseconds)', 'Operations per second', 'Relative (latency)')
    for group, val in sorted(obj.items()):
        buf.write('\n' + '#### ' + group + '\n\n')
        table = []
        for lib in LIBRARIES:
            table.append(
                [lib, val[lib]['median'], '%.1f' % val[lib]['ops'], 0]
            )
        baseline = table[0][1]
        for each in table:
            each[3] = '%.2f' % (each[1] / baseline)
            each[1] = '%.2f' % each[1]
        buf.write(tabulate(table, headers, tablefmt='grid') + '\n')


    print(
        buf.getvalue()
        .replace('-', '')
        .replace('=', '-')
        .replace('+', '|')
        .replace('|||||', '')
        .replace('\n\n', '\n')
    )

try:
    locals()[sys.argv[1]](aggregate())
except KeyError:
    sys.stderr.write("usage: graph (box|tab)\n")
    sys.exit(1)