shapes

#!/usr/bin/env python
# python imports
import os
# rootpy/ROOT imports
from rootpy.extern.argparse import ArgumentParser
from rootpy.plotting import Hist, Graph, Legend
from rootpy.utils.path import mkdir_p
from ROOT import TLatex
# local imports
from mva.samples import (
    Data, QCD, Embedded_Ztautau, MC_Ztautau,
    Pythia_Ztautau, MC_Embedded_Ztautau, Higgs)
from mva.categories import CATEGORIES, Category_Preselection
from mva.variables import VARIABLES
from mva.analysis import Analysis
from mva.templates import RatioPlot
from mva.systematics import iter_systematics
from mva import save_canvas,log
from mva.defaults import TARGET_REGION

parser = ArgumentParser()
parser.add_argument('--year', type=int, default=2012, choices=(2011, 2012))
parser.add_argument('--output-formats', default=['png'], nargs='+',
                    choices=('png', 'eps', 'pdf'),
                    help='output formats')
parser.add_argument('--categories', default='mva_all',
                    choices=CATEGORIES.keys(),
                    help='category definitions')
parser.add_argument('shapes', nargs='*')
args = parser.parse_args()

categories = CATEGORIES[args.categories]+[Category_Preselection]

# list of fields to plot
fields = [
    'mmc1_mass',
    'MET_et',
    'MET_centrality',
    'resonance_pt',
    'dEta_tau1_tau2',
    'dR_tau1_tau2',
    'dPhi_tau1_tau2',
    'tau1_pt',
    'tau2_pt',
    'tau1_eta',
    'tau2_eta',
    'numJets',
]
field_dict = dict([(field, VARIABLES[field]) for field in fields])


def draw_ratio(a, b, field, category, textsize=22):
    plot = RatioPlot(
        xtitle=field_dict[field]['root'],
        ytitle='Normalized Events',
        ratio_title='A / B',
        ratio_range=(0, 2),
        ratio_line_values=[0.5, 1, 1.5])
    a_integral = a.integral()
    if a_integral != 0:
        a /= a_integral
    b_integral = b.integral()
    if b_integral != 0:
        b /= b_integral
    a.title = 'A: ' + a.title
    b.title = 'B: ' + b.title
    a.color = 'black'
    b.color = 'red'
    a.legendstyle = 'L'
    b.legendstyle = 'L'
    a.markersize = 0
    b.markersize = 0
    a.linewidth = 2
    b.linewidth = 2
    a.fillstyle = 'hollow'
    b.fillstyle = 'hollow'
    a.linestyle = 'solid'
    b.linestyle = 'dashed'
    a.drawstyle='hist E0'
    b.drawstyle='hist E0'
    plot.draw('main', [a, b], ypadding=(0.3, 0.))
    ratio = Hist.divide(a, b, fill_value=-1)
    ratio.drawstyle = 'hist'
    ratio.color = 'black'
    ratio_band = Graph(ratio, fillstyle='/', fillcolor='black', linewidth=0)
    ratio_band.drawstyle = '20'
    plot.draw('ratio', [ratio_band, ratio])
    with plot.pad('main') as pad:
        # legend
        leg = Legend([a, b],
                     leftmargin=0.25, topmargin=0.1,
                     margin=0.18, textsize=textsize)
        leg.Draw()
        # draw the category label
        label = TLatex(
            pad.GetLeftMargin() + 0.04, 0.87,
            category.label)
        label.SetNDC()
        label.SetTextFont(43)
        label.SetTextSize(textsize)
        label.Draw()
        # show p-value and chi^2
        pvalue = a.Chi2Test(b, 'WW')
        pvalue_label = TLatex(
            pad.GetLeftMargin() + 0.04, 0.8,
            "p-value={0:.2f}".format(pvalue))
        pvalue_label.SetNDC(True)
        pvalue_label.SetTextFont(43)
        pvalue_label.SetTextSize(textsize)
        pvalue_label.Draw()
        chi2 = a.Chi2Test(b, 'WW CHI2/NDF')
        chi2_label = TLatex(
            pad.GetLeftMargin() + 0.04, 0.72,
            "#frac{{#chi^{{2}}}}{{ndf}}={0:.2f}".format(chi2))
        chi2_label.SetNDC(True)
        chi2_label.SetTextFont(43)
        chi2_label.SetTextSize(textsize)
        chi2_label.Draw()
    return plot


def compare(a, b, category, name):
    a_hists, field_scale = a.get_field_hist(field_dict, category)
    b_hists, _ = b.get_field_hist(field_dict, category)
    a.draw_array(a_hists, category, TARGET_REGION, field_scale=field_scale)
    b.draw_array(b_hists, category, TARGET_REGION, field_scale=field_scale)
    for field in fields:
        # draw ratio plot
        a_hist = a_hists[field]
        b_hist = b_hists[field]
        plot = draw_ratio(a_hist, b_hist, field, category)
        for output in args.output_formats:
            save_canvas(plot, 'plots/shapes', '{0}/shape_{0}_{1}_{2}_{3}.{4}'.format(
                name, field, category.name, args.year % 1000, output))


if not args.shapes or 'data' in args.shapes:
    # data
    data = Data(year=args.year, label='TES-shifted Data')
    data_orig = Data(year=args.year, tes_shift=False, label='Original Data')
    for category in categories:
        # compare tes-shifted data with original data
        compare(data, data_orig, category, 'data_tes')


if not args.shapes or 'ztt' in args.shapes:
    # ztt
    ztt_eb = Embedded_Ztautau(year=args.year, label='Embedded Data')
    ztt_mc = MC_Ztautau(year=args.year, label='MC')
    
    # compare with and without spin weight
    ztt_eb_spin = Embedded_Ztautau(year=args.year, label='TauSpinner')
    ztt_eb_nospin = Embedded_Ztautau(year=args.year, label='No TauSpinner', tauspinner=False)

    # compare with and without posterior trigger correction
    ztt_eb_correct = Embedded_Ztautau(year=args.year, label='Corrected')
    ztt_eb_nocorrect = Embedded_Ztautau(year=args.year, label='Uncorrected', posterior_trigger_correction=False)

    # compare 2011 with 2012 embedding
    ztt_eb_11 = Embedded_Ztautau(year=2011, label='7 TeV Embedding')
    ztt_eb_12 = Embedded_Ztautau(year=2012, label='8 TeV Embedding')

    # MC Ztt and MC embedded Ztt
    #ztt_pyth = Pythia_Ztautau(year=2012, label='MC')
    #ztt_mceb = MC_Embedded_Ztautau(year=2012, label='Embedded MC')
    
    for category in categories:
        # compare Ztt shapes
        compare(ztt_eb, ztt_mc, category, 'ztt')
        # compare with and without tau spinner weight
        compare(ztt_eb_spin, ztt_eb_nospin, category, 'ztt_spin')
        # compare with and without posterior trigger correction
        compare(ztt_eb_correct, ztt_eb_nocorrect, category, 'ztt_posterior_correct')
        # compare MC Ztautau with embedded MC Zmumu
        #compare(ztt_pyth, ztt_mceb, category, 'ebztt')
        # compare 7 TeV with 8 TeV embedding
        compare(ztt_eb_12, ztt_eb_11, category, 'ztt_7_v_8')


if not args.shapes or 'qcd' in args.shapes:
    # fakes
    qcd_nos = Analysis(args.year, fakes_region='nOS').normalize(Category_Preselection).qcd.decorate(label='Isolated nOS Fakes')
    qcd_nos_nonisol = Analysis(args.year, fakes_region='nOS_NONISOL').normalize(Category_Preselection).qcd.decorate(label='Non-isolated nOS Fakes')
    qcd_nos_double_nonisol = Analysis(args.year, fakes_region='nOS_DOUBLE_NONISOL').normalize(Category_Preselection).qcd.decorate(label='Both Non-isolated nOS Fakes')

    qcd_ss = Analysis(args.year, fakes_region='SS').normalize(Category_Preselection).qcd.decorate(label='Isolated SS Fakes')
    # to get OS fakes, use "best" Ztt fits from nOS_NONISOL and then change the
    # shape region to OS and scale to 1.
    qcd_os = Analysis(args.year, fakes_region='nOS_NONISOL').normalize(Category_Preselection).qcd.decorate(label='Isolated OS Fakes')
    qcd_os.shape_region = 'OS'
    qcd_os.scale = 1.

    qcd_ss_nonisol = Analysis(args.year, fakes_region='SS_NONISOL').normalize(Category_Preselection).qcd.decorate(label='Non-isolated SS Fakes')
    qcd_os_nonisol = Analysis(args.year, fakes_region='OS_NONISOL').normalize(Category_Preselection).qcd.decorate(label='Non-isolated OS Fakes')
    qcd_ss_double_nonisol = Analysis(args.year, fakes_region='SS_DOUBLE_NONISOL').normalize(Category_Preselection).qcd.decorate(label='Both Non-isolated SS Fakes')
    qcd_os_double_nonisol = Analysis(args.year, fakes_region='OS_DOUBLE_NONISOL').normalize(Category_Preselection).qcd.decorate(label='Both Non-isolated OS Fakes')

    qcd_nonisol = Analysis(args.year, fakes_region='NONISOL').normalize(Category_Preselection).qcd.decorate(label='Non-isolated Fakes')
    qcd_double_nonisol = Analysis(args.year, fakes_region='DOUBLE_NONISOL').normalize(Category_Preselection).qcd.decorate(label='Both Non-isolated Fakes')


    for category in categories:
        # compare SS and nOS QCD shapes
        compare(qcd_nos, qcd_ss, category, 'qcd_ss_v_nos')

        # compare SS ISOL and SS NONISOL QCD shapes
        compare(qcd_ss, qcd_ss_nonisol, category, 'qcd_ss')
        # compare OS ISOL and OS NONISOL QCD shapes
        compare(qcd_os, qcd_os_nonisol, category, 'qcd_os')
        # compare nOS ISOL and nOS NONISOL QCD shapes
        compare(qcd_nos, qcd_nos_nonisol, category, 'qcd_nos')

        # compare SS NONISOL and OS NONISOL QCD shapes
        compare(qcd_ss_nonisol, qcd_os_nonisol, category, 'qcd_ss_v_os_nonisol')
        # compare SS ISOL and OS ISOL QCD shapes
        compare(qcd_ss, qcd_os, category, 'qcd_ss_v_os_isol')

        # compare nOS NONISOL and OS NONISOL QCD shapes
        compare(qcd_nos_nonisol, qcd_os_nonisol, category, 'qcd_nos_v_os_nonisol')
        # compare nOS ISOL and OS ISOL QCD shapes
        compare(qcd_nos, qcd_os, category, 'qcd_nos_v_os_isol')
        # compare nOS NONISOL and OS ISOL QCD shapes
        compare(qcd_nos_nonisol, qcd_os, category, 'qcd_nos_nonisol_v_os_isol')
        # compare NONISOL and OS ISOL QCD shapes
        compare(qcd_nonisol, qcd_os, category, 'qcd_nonisol_v_os_isol')

        # compare nonisol with double nonisol
        #compare(qcd_os_nonisol, qcd_os_double_nonisol, category, 'qcd_double_os')
        #compare(qcd_ss_nonisol, qcd_ss_double_nonisol, category, 'qcd_double_ss')
        #compare(qcd_nos_nonisol, qcd_nos_double_nonisol, category, 'qcd_double_nos')
        #compare(qcd_nonisol, qcd_double_nonisol, category, 'qcd_double')

        # compare nOS NONISOL and NONISOL (no charge requirement)
        compare(qcd_nos_nonisol, qcd_nonisol, category, 'qcd_nonisol')


if not args.shapes or 'qcd_ztautau' in args.shapes:
    ztt_eb = Embedded_Ztautau(year=args.year, label='Z#rightarrow#tau#tau')
    qcd = Analysis(args.year, fakes_region='nOS_NONISOL').normalize(Category_Preselection).qcd.decorate(label='Fakes')
    higgs = Higgs(args.year, mass=125)
    for category in categories:
        compare(ztt_eb, qcd, category, 'qcd_vs_ztautau')
        compare(ztt_eb, higgs, category, 'higgs_vs_ztautau')