import ROOT
from math import *

###%jsroot on

# WW Background
##bkg = ROOT.TFile.Open("https://atlas-opendata.web.cern.ch/atlas-opendata/samples/2020/2lep/MC/mc_363492.llvv.2lep.root")
bkg = ROOT.TFile.Open("./mc_363492.llvv.2lep.root")
t_bkg = bkg.Get("mini")
n_bkg = t_bkg.GetEntries()
kf = 0.01
n_bkgf = int(n_bkg*kf)

## SM H->WW Signal
##sig = ROOT.TFile.Open("https://atlas-opendata.web.cern.ch/atlas-opendata/samples/2020/2lep/MC/mc_345324.ggH125_WW2lep.2lep.root")
sig = ROOT.TFile.Open("./mc_345324.ggH125_WW2lep.2lep.root")
t_sig = sig.Get("mini")
n_sig = t_sig.GetEntries()
n_sigf = int(n_sig*kf)

c = ROOT.TCanvas("testCanvas","a first way to plot a variable",800,600)
c.Divide(2,2)

h_bgs = ROOT.TH1F("h_bgs","Missing transverse energy",20,0,200)
h2_bgs = ROOT.TH1F("h2_bgs","Number of Leptons",20,0,200)
h3_bgs = ROOT.TH1F("h3_bgs","Transverse mass",20,0,200)

h_sig = ROOT.TH1F("h_sig","Missing transverse energy",20,0,200)
h2_sig = ROOT.TH1F("h2_sig","Number of Leptons",20,0,200)
h3_sig = ROOT.TH1F("h3_sig","Transverse mass",20,0,200)

print("Background sample:")
for n in range(n_bkgf):
    t_bkg.GetEntry(n)
    
    ## printing the evolution in number of events
    if(n%10000==0):
        if n != 0: print("\t",n)
    
    lep_pt = t_bkg.lep_pt.at(0)
    lep_phi = t_bkg.lep_phi.at(0)
    met_et = t_bkg.met_et
    met_phi = t_bkg.met_phi
        
    h_bgs.Fill(met_et/1000.)
    h2_bgs.Fill(lep_pt/1000.)

    mtb = sqrt(2.0*lep_pt*met_et*(1.0-cos(lep_phi-met_phi)))/1000
    h3_bgs.Fill(mtb)

print("Signal sample:")
for m in range(n_sigf):
    t_sig.GetEntry(m)

    ## printing the evolution in number of events
    if(m%2000==0):
        if m != 0: print("\t",m)
    
    lep_pt = t_sig.lep_pt.at(0)
    lep_phi = t_sig.lep_phi.at(0)
    met_et = t_sig.met_et
    met_phi = t_sig.met_phi    
    
    h_sig.Fill(met_et/1000.)
    h2_sig.Fill(lep_pt/1000.)
    
    mts = sqrt(2.0*lep_pt*met_et*(1.0-cos(lep_phi-met_phi)))/1000
    h3_sig.Fill(mts)

ROOT.gStyle.SetOptStat("")    

c.cd(1)
scale_bgs = h_bgs.Integral()
h_bgs.Scale(1/scale_bgs)

scale_sig = h_sig.Integral()
h_sig.Scale(1/scale_sig)

h_bgs.SetFillStyle(3001)
h_bgs.SetFillColor(4)
h_bgs.SetLineColor(4)

h_sig.SetFillStyle(3003)
h_sig.SetFillColor(2)
h_sig.SetLineColor(2)

h_sig.Draw("hist")
h_bgs.Draw("hist,same")

legend=ROOT.TLegend(0.5,0.7,0.9,0.9)
legend.AddEntry(h_bgs,"Background (WW) ","l")
legend.AddEntry(h_sig,"Signal (H #rightarrow WW)","l")

legend.Draw()

c.cd(2)

scale2_bgs = h2_bgs.Integral()
h2_bgs.Scale(1/scale2_bgs)

scale2_sig = h2_sig.Integral()
h2_sig.Scale(1/scale2_sig)

h2_bgs.SetLineColor(4)
h2_sig.SetLineColor(2)

h2_sig.Draw("hist")
h2_bgs.Draw("hist,same")

legend2=ROOT.TLegend(0.5,0.7,0.9,0.9)
legend2.AddEntry(h2_bgs,"Background (WW) ","l")
legend2.AddEntry(h2_sig,"Signal (H #rightarrow WW)","l")
legend2.Draw()

c.cd(3)

scale3_bgs = h3_bgs.Integral()
h3_bgs.Scale(1/scale3_bgs)

scale3_sig = h3_sig.Integral()
h3_sig.Scale(1/scale3_sig)

h3_bgs.SetLineColor(4)
h3_sig.SetLineColor(2)

h3_sig.Draw("hist")
h3_bgs.Draw("hist,same")

legend3=ROOT.TLegend(0.5,0.7,0.9,0.9)
legend3.AddEntry(h3_bgs,"Background (WW) ","l")
legend3.AddEntry(h3_sig,"Signal (H #rightarrow WW)","l")
legend3.Draw()

c.Draw()
c.Update()

c.SaveAs("example1.png")
c.SaveAs("example1.pdf")