# Let's start with importing the needed modules
from ROOT import gStyle, TFile, TH1, TH1D, TH2D, TCanvas, TLegend, TColor, TLatex, SetOwnership
import random
# Now let's set some ROOT styling parameters:
# You do not need to know what they mean, but can directly use these settings

gStyle.SetOptStat(0)
gStyle.SetPalette(1)

gStyle.SetTextFont(42)

gStyle.SetTitleStyle(0000)
gStyle.SetTitleBorderSize(0)
gStyle.SetTitleFont(42)
gStyle.SetTitleFontSize(0.055)

gStyle.SetTitleFont(42, "xyz")
gStyle.SetTitleSize(0.5, "xyz")
gStyle.SetLabelFont(42, "xyz")
gStyle.SetLabelSize(0.45, "xyz")

# A function to apply a fix to the cutflow histograms from CutLang.  
# It takes out the redundant bins showing histogramming.
def fixcutflow(h):
    nbins = 0
    for i in range(1, h.GetXaxis().GetNbins()+1):
        if not ("Histo" in h.GetXaxis().GetBinLabel(i) or "ALL" in h.GetXaxis().GetBinLabel(i)):
            nbins = nbins + 1
    hnew = TH1D(h.GetName()+str(random.random()), h.GetTitle(), nbins, 0, nbins)
    ibin = 0
    for i in range(1, h.GetXaxis().GetNbins()+1):
        label = h.GetXaxis().GetBinLabel(i)
        if not ("Histo" in label or "ALL" in label):
            ibin = ibin + 1
            hnew.GetXaxis().SetBinLabel(ibin, label)
            hnew.SetBinContent(ibin, h.GetBinContent(i))
    h = hnew
    SetOwnership(h, False)
    return h

# Let's open two histoOut files produced by CutLang, whose histograms we will compare. 
# These could, for example. be a signal file and a background file
fs = TFile("../../src/histoOut-ttbartovlq_TT800.root")
fb = TFile("../../src/histoOut-ttbartovlq_SMttbar.root")

# Let's also define some variables necessary for weight calculation
# Cross section
lumi = 2.26
# BG masses : [cross section, all events, total skimmed events]
bgxsec = 831760
bgallevents = 97994442
bgskimevents = 20620375

# Signal masses : [cross section, all events, total skimmed events]
massxsec = {'700' : [455, 805000, 256942],
            '800' : [196, 788200, 252828],
            '1200' : [11.8, 817800, 253993],
            }
mass = '800'
sgxsec = massxsec[mass][0]
sgallevents = massxsec[mass][1]
sgskimevents = massxsec[mass][2]

# Now let's draw some histograms. 
# We will compare distributions for different variables.
# You can try this with different histograms in different regions.
# Provide the region and histo names in the following list:
histoinfos = [
# ["<regionname>", "<histoname>"]
    ["fourjetoneb", "cutflow"],
    ["fourjetoneb", "hmtop1b"],
    ["fourjetoneb", "h2mtop1b"],
    ["fourjettwob", "cutflow"],
    ["fourjettwob", "hmtop"],
    ["fourjettwob", "h2mtop"],
#    ["fourjettwob", "hpttop"],
#    ["fourjettwob", "h2pttop"],
#    ["fourjettwob", "hnjets"],
#    ["fourjettwob", "hj1pT"],
#    ["fourjettwob", "hj2pT"],
#    ["fourjettwob", "hj3pT"],
#    ["fourjettwob", "hmet"],
#    ["fourjettwob", "hST"],
#    ["fourjettwob", "hnak8"],
#    ["optforvlqreint", "cutflow"],
#    ["optforvlqreint", "hak8j1pT"],
#    ["optforvlqreint", "h2ak8j1pT"],
#    ["optforvlqreint", "hak8j1m"],
#    ["optforvlqreint", "h2ak8j1m"],
#    ["optforvlqreint", "h2mtop2"],
#    ["optforvlqreint", "h2pttop2"],
#    ["optforvlqreint", "hmet2"],
#    ["optforvlqreint", "hST2"],
#    ["optforvlqreint", "h2ak8j1pT2"],
#    ["optforvlqreint", "h2ak8j1m2"],
]

# Make lists for convases and legends. This is just a fix for jupyter root!!!
canvases = []
texts = []
legends = []

# Loop over the histograms and prepare the plots:
for hinfo in histoinfos:
    # In which region would you like to draw? You can change the region name. 
    region = hinfo[0]
    # Which histogram would you like to plot?
    hname = hinfo[1]
    # Get the histograms from the file:
    hsg = fs.Get(region+"/"+hname)
    hbg = fb.Get(region+"/"+hname)

    if hname == "cutflow":
        hsg = fixcutflow(hsg)
        hbg = fixcutflow(hbg)       
        
    hsg.SetName(hsg.GetName()+str(random.random()))
    hbg.SetName(hbg.GetName()+str(random.random()))
    
    # Calculate event weights
    sgw = sgxsec * lumi * (sgskimevents / fs.Get(region+"/cutflow").GetBinContent(1)) / sgallevents
    bgw = bgxsec * lumi * (bgskimevents / fb.Get(region+"/cutflow").GetBinContent(1)) / bgallevents
    # Scale the histograms with weights
    hsg.Scale(sgw)
    hbg.Scale(bgw)    

    if hsg.GetMaximum() > hbg.GetMaximum(): 
        hbg.SetMaximum(hsg.GetMaximum()*1.1)
    hbg.SetMinimum(0)
    
    # Format the histograms: scaling, lines, colors, axes titles, etc..  
    # You do not need to learn the commands here unless you are really curious.
    # Otherwise just execute the cell.
        
    # Histogram style settings:
    hsg.SetLineWidth(2)
    hbg.SetLineWidth(2)

    # Set the colors:
    # Color numbers can be retrived from https://root.cern.ch/doc/master/classTColor.html
    # (check for color wheel)
    hbg.SetFillColor(400-7) # kYellow - 7
    hsg.SetLineColor(600) # kBlue
    hbg.SetLineColor(400+2) # kYellow + 2

    # Titles, labels.
    # It is enough to set such variables ONLY FOR THE FIRST HISTOGRAM YOU WILL DRAW
    # i.e., the one you will call by .Draw().  The rest you will draw by .Draw("same") will only 
    # contribute with the historam curve.
    hbg.SetTitle("")
    hbg.GetXaxis().SetTitle(hsg.GetTitle())
    hbg.GetXaxis().SetTitleOffset(1.25)
    hbg.GetXaxis().SetTitleSize(0.05)
    hbg.GetXaxis().SetLabelSize(0.045)
    hbg.GetXaxis().SetNdivisions(8, 5, 0)
    hbg.GetYaxis().SetTitle("number of events")
    #hbg.GetYaxis().SetTitleOffset(1.4)
    hbg.GetYaxis().SetTitleSize(0.05)
    hbg.GetYaxis().SetLabelSize(0.045)
    
    # Write the region name on the histogram
    t = TLatex(0.60, 0.85, region)
    t.SetTextSize(0.042)
    t.SetTextFont(42)
    t.SetNDC()
    texts.append(t)

    # Make a plot legend
    # Change the entry names to reflect processes you are plotting!
    l = TLegend(0.55, 0.70, 0.88, 0.82)
    l.SetBorderSize(0)
    l.SetFillStyle(0000)
    l.AddEntry(hsg,"TT m=800", "l")
    l.AddEntry(hbg,"tt+jets", "f")
    legends.append(l)
    
    # Now we make a canvas and draw our histograms
    cxwidth = 650
    if hname == "cutflow": cxwidth = 1000
    c = TCanvas("c_"+region+"_"+hname, "c_"+region+"_"+hname, cxwidth, 500)
    c.SetBottomMargin(0.15)
    c.SetLeftMargin(0.15)
    c.SetRightMargin(0.15)
    c.SetLogy(1)
    logy = 1
    if logy == 1:
        hbg.SetMaximum(hbg.GetMaximum()*10)
        hbg.SetMinimum(0.5)
    c.SetLogy(logy)
    hbg.Draw("hist")
    hsg.Draw("hist same")
    c.Draw()
    l.Draw("same")
    t.Draw("same")
    canvases.append(c)
    
# Draw the canvases:
for c in canvases:
    c.Draw()