#!/usr/bin/env python
import os
import logging
import re
import rootpy
from rootpy.plotting.style import set_style
from rootpy.tree import Cut

from brtautau.variables import VARIABLES
from brtautau.samples import Higgs
from brtautau.categories.truth import Category_Preselection
from brtautau.plotting import draw_hists
log = logging.getLogger(os.path.basename(__file__))

set_style('ATLAS', shape='rect')

if __name__ == "__main__":
    from argparse import ArgumentParser
    parser = ArgumentParser()
    parser.add_argument('--var', default=None, help='Specify a particular variable')
    parser.add_argument('--logy', default=False, action='store_true')
    parser.add_argument('--cut', default=None, type=str, help='additional cut to apply')
    parser.add_argument('--mode', default='VBF', type=str, choices=['VBF', 'gg'])

    args = parser.parse_args()

    cat = Category_Preselection
    cuts = None
    if args.cut:
        cuts = Cut(args.cut)

    vars = VARIABLES
    if args.var is not None:
        if args.var in VARIABLES:
            vars = {args.var: VARIABLES[args.var]}
        else:
            raise RuntimeError('Wrong key (see variables.py)')


    higgs_100 = Higgs(mass=100, mode=args.mode)
    higgs_125 = Higgs(mass=125, mode=args.mode)
    higgs_150 = Higgs(mass=150, mode=args.mode)

    field_hists = higgs_100.get_field_hist(vars)
    hists_100 = higgs_100.get_hist_array(field_hists, cat, cuts=args.cut)
    hists_125 = higgs_125.get_hist_array(field_hists, cat, cuts=args.cut)
    hists_150 = higgs_150.get_hist_array(field_hists, cat, cuts=args.cut)

    for field, _ in field_hists.items():
        h_100 = hists_100[field]
        h_125 = hists_125[field]
        h_150 = hists_150[field]

        plot = draw_hists([h_100, h_125, h_150], field, cat, logy=args.logy, unit_area=True)
        log.info(plot)
        if args.logy:
            plot.SaveAs(os.path.join('plots', 'lineshape_'+field+'_'+cat.name+'_logy.png'))
        else: 
            plot.SaveAs(os.path.join('plots', 'lineshape_'+field+'_'+cat.name+'.png'))

                               
