from collections import OrderedDict
import subprocess
import os

def run_perf(args, json_name):
    if os.path.exists(json_name):
        return
    args = args + ['-o', json_name]
    print(args)
    p = subprocess.Popen(
        args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    out, err = p.communicate()
    print(out.decode())
    print(err.decode())


setup = OrderedDict([
    ('astropy', 'import astropy.units as u'),
    ('pint', 'from pint import UnitRegistry; u = UnitRegistry()'),
    ('unyt', 'import unyt as u'),
    ('yt', 'import yt.units as u'),
    ('numpy', 'import numpy as np'),
])

base_args = ['python3', '-m', 'perf', 'timeit', '--rigorous']

shared_setup = 'import numpy as np; import operator'

base_setups = OrderedDict([
    ('small_list', 'data = [1., 2., 3.]'),
    ('small_tuple', 'data = (1., 2., 3.)'),
    ('small_array', 'data = np.array([1., 2., 3.])'),
    ('big_list', 'data = np.arange(1e6).tolist()'),
    ('big_array', 'data = np.arange(1e6)'),
])

op_ufuncs = OrderedDict([
    ('operator.add', 'np.add'),
    ('operator.sub', 'np.subtract'),
    ('operator.mul', 'np.multiply'),
    ('operator.truediv', 'np.true_divide'),
])

for bs in base_setups:
    for package in sorted(setup):
        print(package)
        setup_s = '; '.join([shared_setup, setup[package], base_setups[bs]])
        args = base_args + ['-s', setup_s + ' ']
        if package == 'numpy':
            args.append('np.asarray(data)')
        else:
            args.append('data*u.g')
        json_name = '{}_{}_create.json'.format(package, bs)
        run_perf(args, json_name)

for bs in base_setups:
    for op, ufunc in op_ufuncs.items():
        for bench, bench_name in [
                (op + r'(data1, data2)', op + '12.json'),
                (op + r'(data2, data1)', op + '21.json'),
                (ufunc + r'(data1, data2)', ufunc + '12.json'),
                (ufunc + r'(data2, data1)', ufunc + '21.json'),
                (ufunc + r'(data1, data2, out=out)', ufunc + '12out.json'),
                (ufunc + r'(data2, data1, out=out)', ufunc + '21out.json'),
        ]:
            for unit_choice in [('g', 'g'), ('kg', 'g')]:
                for package in sorted(setup):
                    print(package)
                    setup_s = '; '.join(
                        [shared_setup, setup[package], base_setups[bs]]) + '; '
                    if 'out' in bench:
                        if package not in ('pint', 'numpy'):
                            setup_s += 'out=data*u.{}; '.format(unit_choice[0])
                        else:
                            setup_s += 'out=np.array(data); '
                    if package == 'numpy':
                        setup_s += '; '.join([r'data1 = np.array(data)',
                                              r'data2 = np.array(data)'])
                        if unit_choice[0] != unit_choice[1]:
                            _bench = bench.replace('data1', '.001*data1')
                    else:
                        setup_s += '; '.join(
                            ['data1 = data*u.{}'.format(unit_choice[0]),
                             'data2 = data*u.{}'.format(unit_choice[1])])
                        _bench = bench
                    args = base_args + ['-s', setup_s + ' ']
                    json_name = '{}_{}_{}{}'.format(
                        package, bs, unit_choice[0], unit_choice[1])
                    run_perf(args + [_bench], json_name + '_' + bench_name)
    

@@ -1,39 +1,150 @@
-from collections import OrderedDict
+import matplotlib%0Amatplotlib.use('agg')%0Afrom collections import OrderedDict%0Afrom matplotlib import pyplot as plt%0Aimport json%0Aimport os%0Aimport perf
 %0Aimp
@@ -157,26 +157,16 @@
 bprocess
-%0Aimport os
 %0A%0Adef ru
@@ -460,21 +460,51 @@
 de())%0A%0A%0A
-setup
+def make_plot(extension):%0A    means
  = Order
@@ -514,52 +514,639 @@
 ict(
-%5B
+)
 %0A    
-('astropy', 'import astropy.units as u
+stddevs = OrderedDict()%0A    for package in setup:%0A        benchmark = perf.Benchmark.load(%0A            open('%7B%7D_%7B%7D'.format(package, extension), 'r'))%0A        means%5Bpackage%5D = benchmark.mean()%0A        stddevs%5Bpackage%5D = benchmark.stdev()%0A    fig, ax = plt.subplots()%0A    packages = means.keys()%0A    means = means.values()%0A    stddevs = stddevs.values()%0A    ax.bar(packages, means, yerr=stddevs)%0A    fig.suptitle(extension.replace('.json', '').replace('_', ' ').title())%0A    ax.set_ylabel('time (s); lower is better')%0A    plt.savefig(extension.replace('.json', '.png'))%0A%0A%0Asetup = OrderedDict(%5B%0A    ('numpy', 'import numpy as np
 '),%0A
@@ -1218,20 +1218,23 @@
 ,%0A    ('
-unyt
+astropy
 ', 'impo
@@ -1240,42 +1240,15 @@
 ort 
-unyt as u'),%0A    ('yt', 'import yt
+astropy
 .uni
@@ -1260,29 +1260,28 @@
  u'),%0A    ('
-numpy
+unyt
 ', 'import n
@@ -1275,35 +1275,33 @@
 t', 'import 
-numpy
+unyt
  as 
-np
+u
 '),%0A%5D)%0A%0Abase
@@ -2286,16 +2286,59 @@
 on_name)
+%0A    make_plot(%22%7B%7D_create.json%22.format(bs))
 %0A%0Afor bs
@@ -4206,17 +4206,133 @@
 bench_name)%0A
+                make_plot(%22%7B%7D_%7B%7D%7B%7D_%7B%7D%22.format(%0A                    bs, unit_choice%5B0%5D, unit_choice%5B1%5D, bench_name))%0A
     %0A

make plots for benchmarks, drop yt