#!/usr/bin/env python

from root_numpy import root2array

import argparse
import itertools
import matplotlib
matplotlib.use('Agg')
import matplotlib.gridspec as gs
import matplotlib.pyplot as plt
import numpy as np
import os
import pandas as pd


matplotlib.rc('axes.formatter', limits=(-3, 4))


html = """
<hmtl>
<head>
<style>
tr:first-child {{background-color: #ccc}}
tr:nth-child(even) {{background-color: #f2f2f2}}
table {{font-size: 80%}}
.highlight {{color: #e00; font-weight: bold}};
</style>
<title>Comparing {b} and {c}</title>
</head>
<body>
<section>
<h1>Event Differences</h1>
{}
<h1>Comparison Tables</h1>
{}
</section>
<section>
<h1>Comparison Plots</h1>
{}
</section>
</body>
</html>
"""

img = '<a href="{0}.png"><img src="{0}.png" width="600" height="400"></img></a>'
row = '<tr>{}</tr>'
cell = '<td class="{0[1]}">{0[0]}</td>'
table = '<h2>{}</h2>{}<table>{}</table>'
line = '<p>{}</p>'

absolute_precision = 0.001
relative_precision = 0.02


def bin_it(df, mn, mx):
    bins = np.linspace(mn, mx, 30)
    groups = df.groupby(pd.cut(df, bins))
    return bins, groups


def load_root(spec):
    fn, tn, label = spec.split(':')
    return label, root2array(fn, tn)


def load_txt(spec, names):
    fn, label = spec.split(':')
    return label, np.genfromtxt(fn, delimiter=' ', names=(True if not names else names))


def load(spec, names):
    if ".root" in spec:
        return load_root(spec)
    else:
        return load_txt(spec, names)


def compare(joint, quants, limit):
    def highlight(a, b):
        return a != b and \
            abs(a - b) > relative_precision * a and \
            abs(a - b) > absolute_precision

    cols = ['nEvent', 'rank'] + ['_'.join(t) for t in itertools.product(quants, [basename, compname])]
    rows = [''.join(map(cell.format, [(c, '') for c in cols]))]
    count = 0
    for index, series in joint.iterrows():
        values_b = np.array([series[q + '_b'] for q in quants])
        values_c = np.array([series[q + '_c'] for q in quants])

        if not np.any(
                np.logical_and(
                    np.logical_and(
                        np.absolute(values_b - values_c) > relative_precision * np.absolute(values_b),
                        np.absolute(values_b - values_c) > absolute_precision
                    ),
                    values_b != values_c
                )):
            continue

        count += 1
        if count > limit and limit > 0:
            break

        rows += [''.join(map(cell.format, [(series['nEvent'], ''), (index, '')] + [(v, '' if not highlight(t[0], t[1]) else 'highlight') for t in zip(values_b, values_c) for v in t]))]
    return rows


def valid(data):
    if len(data) == 0:
        return False
    uniq = np.unique(data)
    if len(uniq) == 1 and uniq[0] < -99.:
        return False
    return True


def compare_quantities(joint, names, limit):
    quantify = line.format('The first {} differences'.format(limit) if limit > 0 else 'All differences')
    if names:
        rows = compare(joint, names[1:], limit)
        if len(rows) > 0:
            yield ('Differences', quantify, ''.join(map(row.format, rows)))
    else:
        objs = ['ele0', 'ele1', 'mu0', 'mu1', 'tau0', 'tau1', 'jet0', 'jet1', 'jet2', 'jet3']
        for obj in objs:
            quants = sorted([col for col in both if col.startswith(obj)])
            rows = compare(joint, quants, limit)
            if len(rows) > 0:
                yield ('Differences in {}'.format(obj), quantify, ''.join(map(row.format, rows)))
        quants = sorted(
            [col for col in both if col.endswith('weight') or col.startswith('n_')] +
            ['ntags_loose', 'avg_dr_jet', 'tt_deltaR', 'tt_mvis', 'HT', 'MVA_1l2tau_ttbar']
        )
        quants = [q for q in quants if q + '_b' in list(joint) and q + '_c' in list(joint)]
        rows = compare(joint, quants, limit)
        if len(rows) > 0:
            yield ('Differences in events', quantify, ''.join(map(row.format, rows)))


def compare_events(joint, basename, baseframe, compname, compframe):
    only_base = baseframe[~baseframe['nEvent'].isin(joint['nEvent'])]
    only_comp = compframe[~compframe['nEvent'].isin(joint['nEvent'])]

    print "# of events only in baseline:   {}".format(len(only_base))
    print "# of events only in comparison: {}".format(len(only_comp))
    print "# of events shared:             {}".format(len(joint))

    e_base = 'Events in {}: {}'.format(basename, ', '.join([str(n) for n in only_base.nEvent]))
    e_comp = 'Events in {}: {}'.format(compname, ', '.join([str(n) for n in only_comp.nEvent]))

    voi = ['nEvent'] + \
        ['n_fakeablesel_' + o for o in ('mu', 'ele')] + \
        ['n_presel_' + o for o in ('tau', 'jet')]
    voi += ['{}_{}'.format(o, v) for o, v in itertools.product(('mu0', 'ele0'), ('pt', 'eta', 'phi', 'jetCSV', 'leptonMVA'))]
    voi += ['{}_{}'.format(o, v) for o, v in itertools.product(('tau0', 'tau1'), ('pt', 'eta', 'phi', 'charge'))]

    rows = [''.join([cell.format((h, '')) for h in voi])]
    for _, e in only_base.iterrows():
        rows.append(''.join([cell.format((e[v], '')) for v in voi]))
    yield ('Only in ' + basename, line.format(e_base), ''.join([row.format(r) for r in rows]))

    rows = [''.join([cell.format((h, '')) for h in voi])]
    for _, e in only_comp.iterrows():
        rows.append(''.join([cell.format((e[v], '')) for v in voi]))
    yield ('Only in ' + compname, line.format(e_comp), ''.join([row.format(r) for r in rows]))

parser = argparse.ArgumentParser()
parser.add_argument('-n', metavar='EVENTS', default=10, type=int, help='number of events to limit comparison table to')
parser.add_argument('-p', default=False, type=bool, help='produce comparison plots')
parser.add_argument('--columnnames', metavar='FILE', help='column names, one per line')
parser.add_argument(dest='outdir', help='directory to store output in')
parser.add_argument(dest='baseline', help='either: rootfile:treename:label | textfile:label')
parser.add_argument(dest='comparison', help='either: rootfile:treename:label | textfile:label')
args = parser.parse_args()

if not os.path.exists(args.outdir):
    os.makedirs(args.outdir)

names = None
if args.columnnames:
    with open(args.columnnames) as fd:
        names = map(str.strip, fd.readlines())

basename, basetree = load(args.baseline, names)
compname, comptree = load(args.comparison, names)

baseonly = set(basetree.dtype.names) - set(comptree.dtype.names)
componly = set(comptree.dtype.names) - set(basetree.dtype.names)
both = set(comptree.dtype.names) & set(basetree.dtype.names)

print "Fields only in baseline: {}".format(", ".join(sorted(baseonly)))
print "Fields only in comparison: {}".format(", ".join(sorted(componly)))

baseframe = pd.DataFrame(basetree)
compframe = pd.DataFrame(comptree)

joint = pd.merge(baseframe, compframe, how='inner', on=['nEvent'], suffixes=('_b', '_c'))

qcomparison = compare_quantities(joint, names, args.n)
ecomparison = compare_events(joint, basename, baseframe, compname, compframe)

with open(os.path.join(args.outdir, 'index.html'), 'w') as f:
    f.write(
        html.format(
            '\n'.join([table.format(*t) for t in ecomparison]),
            '\n'.join([table.format(*t) for t in qcomparison]),
            '' if not args.p else '\n'.join([img.format(c) for c in sorted(both) if c != 'nEvent']),
            b=basename, c=compname
        ))

if args.p:
    for column in both:
        if column == 'nEvent':
            continue
        print "Processing {}".format(column)
        data_b = joint[column + '_b']
        data_c = joint[column + '_c']

        fig = plt.figure(figsize=(15, 10))

        g1 = gs.GridSpec(3, 1, height_ratios=[2, 2, 1])
        g1.update(top=0.95, bottom=0.05, left=0.05, right=0.59, hspace=0)

        ax1 = plt.subplot(g1[0, 0])
        ax2 = plt.subplot(g1[1, 0], sharex=ax1)
        ax3 = plt.subplot(g1[2, 0], sharex=ax2)

        ax1.set_title(column)
        ax3.set_ylabel('({c} - {b}) / {b}'.format(c=compname, b=basename))

        plt.setp(ax1.get_xticklabels(), visible=False)
        plt.setp(ax2.get_xticklabels(), visible=False)

        g2 = gs.GridSpec(2, 1)
        g2.update(top=0.95, bottom=0.05, left=0.61, right=0.95, hspace=0)

        ax4 = plt.subplot(g2[0, 0])
        ax5 = plt.subplot(g2[1, 0], sharex=ax4)

        ax4.set_title('{c} - {b}'.format(c=compname, b=basename))
        ax4.yaxis.tick_right()
        ax4.yaxis.set_ticks_position('both')
        ax5.yaxis.tick_right()
        ax5.yaxis.set_ticks_position('both')

        plt.setp(ax4.get_xticklabels(), visible=False)

        diff = (data_c - data_b)[(data_b > -666) & (data_c > -666)]

        data_b = data_b[data_b > -666]
        data_c = data_c[data_c > -666]

        mn = -11
        mx = 0
        if len(data_b) > 0:
            mn = min(data_b.min(), mn)
            mx = max(data_b.max(), mx)
        if len(data_c) > 0:
            mn = min(data_c.min(), mn)
            mx = max(data_c.max(), mx)

        if len(diff) > 0:
            try:
                ax4.hist(diff, bins=20)
                ax5.hist(diff, bins=20)
            except KeyError:
                pass
            except AttributeError:
                pass

        if not (valid(data_b) or valid(data_c)):
            print "Skipping!"
            continue

        try:
            bins, groups_b = bin_it(data_b, mn, mx)
            centers = [.5 * (a + b) for a, b in zip(bins, bins[1:])]
            ax1.errorbar(x=centers, y=groups_b.count(), yerr=groups_b.std(), fmt='o', label=basename)
            ax2.errorbar(x=centers, y=groups_b.count(), yerr=groups_b.std(), fmt='o')

            bins, groups_c = bin_it(data_c, mn, mx)
            centers = [.5 * (a + b) for a, b in zip(bins, bins[1:])]
            ax1.errorbar(x=centers, y=groups_c.count(), yerr=groups_c.std(), fmt='o', label=compname)
            ax2.errorbar(x=centers, y=groups_c.count(), yerr=groups_c.std(), fmt='o')

            ax3.errorbar(x=centers, y=groups_b.count() / groups_c.count(), fmt='o')
        except ValueError as e:
            pass

        ax1.legend()
        ax1.set_ylim(ymin=0)
        # ax2.set_ylim(ymin=0)

        if len(data_b) > 0 or len(data_c) > 0:
            try:
                ax2.set_yscale('log')
            except ValueError as e:
                pass
        if len(diff) > 0:
            try:
                ax5.set_yscale('log', nonposy='clip')
            except ValueError as e:
                pass

        plt.savefig(os.path.join(args.outdir, column + '.png'))
        plt.close()

cf02841420b62b771a8179397823e551023e63a3
#!/usr/bin/env python

from root_numpy import root2array

import argparse
import itertools
import matplotlib
matplotlib.use('Agg')
import matplotlib.gridspec as gs
import matplotlib.pyplot as plt
import numpy as np
import os
import pandas as pd


matplotlib.rc('axes.formatter', limits=(-3, 4))


html = """
<hmtl>
<head>
<style>
tr:first-child {{background-color: #ccc}}
tr:nth-child(even) {{background-color: #f2f2f2}}
table {{font-size: 80%}}
.highlight {{color: #e00; font-weight: bold}};
</style>
<title>Comparing {b} and {c}</title>
</head>
<body>
<section>
<h1>Event Differences</h1>
{}
<h1>Comparison Tables</h1>
{}
</section>
<section>
<h1>Comparison Plots</h1>
{}
</section>
</body>
</html>
"""

img = '<a href="{0}.png"><img src="{0}.png" width="600" height="400"></img></a>'
row = '<tr>{}</tr>'
cell = '<td class="{0[1]}">{0[0]}</td>'
table = '<h2>{}</h2>{}<table>{}</table>'
line = '<p>{}</p>'

absolute_precision = 0.001
relative_precision = 0.02


def bin_it(df, mn, mx):
    bins = np.linspace(mn, mx, 30)
    groups = df.groupby(pd.cut(df, bins))
    return bins, groups


def load_root(spec):
    fn, tn, label = spec.split(':')
    return label, root2array(fn, tn)


def load_txt(spec, names):
    fn, label = spec.split(':')
    return label, np.genfromtxt(fn, delimiter=' ', names=(True if not names else names))


def load(spec, names):
    if ".root" in spec:
        return load_root(spec)
    else:
        return load_txt(spec, names)


def compare(joint, quants, limit):
    def highlight(a, b):
        return a != b and \
            abs(a - b) > relative_precision * a and \
            abs(a - b) > absolute_precision

    cols = ['nEvent', 'rank'] + ['_'.join(t) for t in itertools.product(quants, [basename, compname])]
    rows = [''.join(map(cell.format, [(c, '') for c in cols]))]
    count = 0
    for index, series in joint.iterrows():
        values_b = np.array([series[q + '_b'] for q in quants])
        values_c = np.array([series[q + '_c'] for q in quants])

        if not np.any(
                np.logical_and(
                    np.logical_and(
                        np.absolute(values_b - values_c) > relative_precision * np.absolute(values_b),
                        np.absolute(values_b - values_c) > absolute_precision
                    ),
                    values_b != values_c
                )):
            continue

        count += 1
        if count > limit and limit > 0:
            break

        rows += [''.join(map(cell.format, [(series['nEvent'], ''), (index, '')] + [(v, '' if not highlight(t[0], t[1]) else 'highlight') for t in zip(values_b, values_c) for v in t]))]
    return rows


def valid(data):
    if len(data) == 0:
        return False
    uniq = np.unique(data)
    if len(uniq) == 1 and uniq[0] < -99.:
        return False
    return True


def compare_quantities(joint, names, limit):
    quantify = line.format('The first {} differences'.format(limit) if limit > 0 else 'All differences')
    if names:
        rows = compare(joint, names[1:], limit)
        if len(rows) > 0:
            yield ('Differences', quantify, ''.join(map(row.format, rows)))
    else:
        objs = ['ele0', 'ele1', 'mu0', 'mu1', 'tau0', 'tau1', 'jet0', 'jet1', 'jet2', 'jet3']
        for obj in objs:
            quants = sorted([col for col in both if col.startswith(obj)])
            rows = compare(joint, quants, limit)
            if len(rows) > 0:
                yield ('Differences in {}'.format(obj), quantify, ''.join(map(row.format, rows)))
        quants = sorted(
            [col for col in both if col.endswith('weight') or col.startswith('n_')] +
            ['ntags_loose', 'avg_dr_jet', 'tt_deltaR', 'tt_mvis', 'HT', 'MVA_1l2tau_ttbar']
        )
        quants = [q for q in quants if q + '_b' in list(joint) and q + '_c' in list(joint)]
        rows = compare(joint, quants, limit)
        if len(rows) > 0:
            yield ('Differences in events', quantify, ''.join(map(row.format, rows)))


def compare_events(joint, basename, baseframe, compname, compframe):
    only_base = baseframe[~baseframe['nEvent'].isin(joint['nEvent'])]
    only_comp = compframe[~compframe['nEvent'].isin(joint['nEvent'])]

    print "# of events only in baseline:   {}".format(len(only_base))
    print "# of events only in comparison: {}".format(len(only_comp))
    print "# of events shared:             {}".format(len(joint))

    e_base = 'Events in {}: {}'.format(basename, ', '.join([str(n) for n in only_base.nEvent]))
    e_comp = 'Events in {}: {}'.format(compname, ', '.join([str(n) for n in only_comp.nEvent]))

    voi = ['nEvent'] + \
        ['n_fakeablesel_' + o for o in ('mu', 'ele')] + \
        ['n_presel_' + o for o in ('tau', 'jet')]
    voi += ['{}_{}'.format(o, v) for o, v in itertools.product(('mu0', 'ele0'), ('pt', 'eta', 'phi', 'jetCSV', 'leptonMVA'))]
    voi += ['{}_{}'.format(o, v) for o, v in itertools.product(('tau0', 'tau1'), ('pt', 'eta', 'phi', 'charge'))]

    rows = [''.join([cell.format((h, '')) for h in voi])]
    for _, e in only_base.iterrows():
        rows.append(''.join([cell.format((e[v], '')) for v in voi]))
    yield ('Only in ' + basename, line.format(e_base), ''.join([row.format(r) for r in rows]))

    rows = [''.join([cell.format((h, '')) for h in voi])]
    for _, e in only_comp.iterrows():
        rows.append(''.join([cell.format((e[v], '')) for v in voi]))
    yield ('Only in ' + compname, line.format(e_comp), ''.join([row.format(r) for r in rows]))

parser = argparse.ArgumentParser()
parser.add_argument('-n', metavar='EVENTS', default=10, type=int, help='number of events to limit comparison table to')
parser.add_argument('-p', default=False, type=bool, help='produce comparison plots')
parser.add_argument('--columnnames', metavar='FILE', help='column names, one per line')
parser.add_argument(dest='outdir', help='directory to store output in')
parser.add_argument(dest='baseline', help='either: rootfile:treename:label | textfile:label')
parser.add_argument(dest='comparison', help='either: rootfile:treename:label | textfile:label')
args = parser.parse_args()

if not os.path.exists(args.outdir):
    os.makedirs(args.outdir)

names = None
if args.columnnames:
    with open(args.columnnames) as fd:
        names = map(str.strip, fd.readlines())

basename, basetree = load(args.baseline, names)
compname, comptree = load(args.comparison, names)

baseonly = set(basetree.dtype.names) - set(comptree.dtype.names)
componly = set(comptree.dtype.names) - set(basetree.dtype.names)
both = set(comptree.dtype.names) & set(basetree.dtype.names)

print "Fields only in baseline: {}".format(", ".join(sorted(baseonly)))
print "Fields only in comparison: {}".format(", ".join(sorted(componly)))

baseframe = pd.DataFrame(basetree)
compframe = pd.DataFrame(comptree)

joint = pd.merge(baseframe, compframe, how='inner', on=['nEvent'], suffixes=('_b', '_c'))

qcomparison = compare_quantities(joint, names, args.n)
ecomparison = compare_events(joint, basename, baseframe, compname, compframe)

with open(os.path.join(args.outdir, 'index.html'), 'w') as f:
    f.write(
        html.format(
            '\n'.join([table.format(*t) for t in ecomparison]),
            '\n'.join([table.format(*t) for t in qcomparison]),
            '' if not args.p else '\n'.join([img.format(c) for c in sorted(both) if c != 'nEvent']),
            b=basename, c=compname
        ))

if args.p:
    for column in both:
        if column == 'nEvent':
            continue
        print "Processing {}".format(column)
        data_b = joint[column + '_b']
        data_c = joint[column + '_c']

        fig = plt.figure(figsize=(15, 10))

        g1 = gs.GridSpec(3, 1, height_ratios=[2, 2, 1])
        g1.update(top=0.95, bottom=0.05, left=0.05, right=0.59, hspace=0)

        ax1 = plt.subplot(g1[0, 0])
        ax2 = plt.subplot(g1[1, 0], sharex=ax1)
        ax3 = plt.subplot(g1[2, 0], sharex=ax2)

        ax1.set_title(column)
        ax3.set_ylabel('({c} - {b}) / {b}'.format(c=compname, b=basename))

        plt.setp(ax1.get_xticklabels(), visible=False)
        plt.setp(ax2.get_xticklabels(), visible=False)

        g2 = gs.GridSpec(2, 1)
        g2.update(top=0.95, bottom=0.05, left=0.61, right=0.95, hspace=0)

        ax4 = plt.subplot(g2[0, 0])
        ax5 = plt.subplot(g2[1, 0], sharex=ax4)

        ax4.set_title('{c} - {b}'.format(c=compname, b=basename))
        ax4.yaxis.tick_right()
        ax4.yaxis.set_ticks_position('both')
        ax5.yaxis.tick_right()
        ax5.yaxis.set_ticks_position('both')

        plt.setp(ax4.get_xticklabels(), visible=False)

        diff = (data_c - data_b)[(data_b > -666) & (data_c > -666)]

        data_b = data_b[data_b > -666]
        data_c = data_c[data_c > -666]

        mn = -11
        mx = 0
        if len(data_b) > 0:
            mn = min(data_b.min(), mn)
            mx = max(data_b.max(), mx)
        if len(data_c) > 0:
            mn = min(data_c.min(), mn)
            mx = max(data_c.max(), mx)

        if len(diff) > 0:
            try:
                ax4.hist(diff, bins=20)
                ax5.hist(diff, bins=20)
            except KeyError:
                pass
            except AttributeError:
                pass

        if not (valid(data_b) or valid(data_c)):
            print "Skipping!"
            continue

        try:
            bins, groups_b = bin_it(data_b, mn, mx)
            centers = [.5 * (a + b) for a, b in zip(bins, bins[1:])]
            ax1.errorbar(x=centers, y=groups_b.count(), yerr=groups_b.std(), fmt='o', label=basename)
            ax2.errorbar(x=centers, y=groups_b.count(), yerr=groups_b.std(), fmt='o')

            bins, groups_c = bin_it(data_c, mn, mx)
            centers = [.5 * (a + b) for a, b in zip(bins, bins[1:])]
            ax1.errorbar(x=centers, y=groups_c.count(), yerr=groups_c.std(), fmt='o', label=compname)
            ax2.errorbar(x=centers, y=groups_c.count(), yerr=groups_c.std(), fmt='o')

            ax3.errorbar(x=centers, y=groups_b.count() / groups_c.count(), fmt='o')
        except ValueError as e:
            pass

        ax1.legend()
        ax1.set_ylim(ymin=0)
        # ax2.set_ylim(ymin=0)

        if len(data_b) > 0 or len(data_c) > 0:
            try:
                ax2.set_yscale('log')
            except ValueError as e:
                pass
        if len(diff) > 0:
            try:
                ax5.set_yscale('log', nonposy='clip')
            except ValueError as e:
                pass

        plt.savefig(os.path.join(args.outdir, column + '.png'))
        plt.close()
