import math
import os, shlex
import subprocess
import glob
import shutil
import json
import ROOT
ROOT.gROOT.SetBatch(True)
from collections import OrderedDict
import CMS_lumi, tdrstyle
from optparse import OptionParser
from array import array
import json
import numpy as np



def plot(filepath, channel, spinCase, mass, outdir):
    W = 800
    H  = 600
    T = 0.08*H
    B = 0.12*H
    L = 0.12*W
    R = 0.04*W
    c = ROOT.TCanvas("c","c",100,100,W,W)
    c.SetFillColor(0)
    c.SetBorderMode(0)
    c.SetFrameFillStyle(0)
    c.SetFrameBorderMode(0)
    c.SetLeftMargin( 1.2*L/W )
    c.SetRightMargin( 1.1*R/W )
    c.SetTopMargin( T/H )
    c.SetBottomMargin( B/H )
    c.SetTickx(0)
    c.SetTicky(0)
    c.cd()
    
    
    CMS_lumi.cmsText = "CMS"
    CMS_lumi.extraText = "       Work in Progress"
    
    CMS_lumi.cmsTextSize = 0.65
    CMS_lumi.outOfFrame = True
    
    #f1 = ROOT.TFile("resExtendRaw/2lss/Run2/newProcName/addSystFakeRates_hh_2lss_leq1tau_hh_2lss_leq1tau_SS_MVAOutput_%s_%s.root"%(mass, spinCase))
    f1 = ROOT.TFile(filepath)
    signal_hwwww = f1.Get("signal_ggf_%s_%s_hwwhww"%(spinCase, mass))
    signal_hwwtt = f1.Get("signal_ggf_%s_%s_htthww"%(spinCase, mass))
    signal_htttt = f1.Get("signal_ggf_%s_%s_htthtt"%(spinCase, mass))
    signal  = ROOT.TH1D("signal", "signal", signal_hwwww.GetNbinsX(),signal_hwwww.GetXaxis().GetXmin(),signal_hwwww.GetXaxis().GetXmax())
    WZ = f1.Get("WZ")
    ZZ = f1.Get("ggZZ")
    ZZ.Add(f1.Get("qqZZ"))
    fakes = f1.Get("multilepton_data_fakes")
    signal.Add(signal_hwwww)
    signal.Add(signal_hwwtt)
    signal.Add(signal_htttt)
    signal.SetLineWidth(3)
    signal.SetLineColor(ROOT.kBlue)
    signal.DrawNormalized("hist")
    signal.SetMaximum(1.7*signal.GetMaximum())
    signal.SetTitle(';BDT output; Normalized Events')
    signal.GetXaxis().SetTitleOffset(1.3)
    signal.GetYaxis().SetTitleOffset(1.7)
    signal.GetYaxis().SetTitleFont(43)
    signal.GetXaxis().SetTitleFont(43)
    signal.GetYaxis().SetLabelSize(30)
    signal.GetYaxis().SetLabelOffset(0.015)
    signal.GetYaxis().SetLabelFont(43)
    signal.GetYaxis().SetTitleSize(35)
    signal.GetXaxis().SetTitleSize(35)
    signal.DrawNormalized("hist")
    WZ.SetLineWidth(3)
    WZ.SetLineColor(ROOT.kRed)
    WZ.DrawNormalized("histsame")
    ZZ.SetLineWidth(3)
    ZZ.SetLineColor(ROOT.kRed)
    ZZ.SetLineStyle(9)
    ZZ.DrawNormalized("histsame")
    fakes.SetLineWidth(3)
    fakes.SetLineStyle(2)
    fakes.SetLineColor(ROOT.kRed)
    fakes.DrawNormalized("histsame")
    legend = ROOT.TLegend(0.2,0.7,0.93,0.88)
    legend.SetNColumns(1)
    legend.SetFillStyle(0)
    legend.SetTextFont(42)
    legend.SetBorderSize(0)
    legend.AddEntry(signal, "HH signal (%s %s GeV)"%(spinCase, mass),"L")
    legend.AddEntry(WZ, "WZ","L")
    legend.AddEntry(ZZ, "ZZ","L")
    legend.AddEntry(fakes, "Misidentified Leptons","L")
    legend.Draw()
    
    latex2 = ROOT.TMathText()
    latex2.SetNDC()
    latex2.SetTextAlign(23)
    latex2.SetTextAngle(0)
    latex2.SetTextSize(1.5*latex2.GetTextSize()*0.8)
    latex2.DrawMathText(0.8,0.8, "%s"%channel)
    
    CMS_lumi.lumi_sqrtS = "138 fb^{-1} (13 TeV)"
    CMS_lumi.CMS_lumi(c,0,0)
    c.SaveAs("%s/%s_%s_%s.png"%(outdir, channel, spinCase, mass))


outdir = "resExtendRaw/plots"
masses = ["800","900","1000","1250","1500", "1750", "2000", "2500", "3000"]
filepath = ""
tdrstyle.setTDRStyle()
for mass in masses:
    for channel in ["3l","2lss"]:
        for spinCase in ["spin0","spin2"]:
            if channel is "3l":
                filepath = "resExtendRaw/3l/Run2/newProcName/addSystFakeRates_hh_3l_hh_3l_OS_MVAOutput_%s_%s.root"%(mass, spinCase)
            else:
                filepath = "resExtendRaw/2lss/Run2/newProcName/addSystFakeRates_hh_2lss_leq1tau_hh_2lss_leq1tau_SS_MVAOutput_%s_%s.root"%(mass, spinCase)
            plot(filepath, channel, spinCase, mass, outdir)
