from ROOT import TFile, RooRealVar, RooDataSet, RooArgSet, RooGaussian, RooChebychev, RooArgList, RooAddPdf, RooFit, TH1F
import rootnotes

inputfile = TFile("TESHEP_Pedagogical_Data.root")
inputchain = inputfile.Get("DecayTree")

#Define the variables to be plotted and their upper and lower ranges
vars = {"D0_PT" 	: [2000.	, 10000.], 
        "D0_DIRA_OWNPV" : [0.999	, 1.	],
	"D0_TAU" 	: [0.00015	, 0.004	],
	"D0_MINIPCHI2"  : [0.   	, 50.	]}


#Define the mass variable
nameofmassvar = "D0_MM"
#Define the signal range (from output of previous script)
sigrangelow  = 1849.
sigrangehigh = 1880.
#Define the sideband range to be used for the background extraction
backgroundrangelow  = 1890.
backgroundrangehigh = 1914.
#Define the background scale factor (from output of previous script)
backscale = 1.33


#Define the histograms -- for each variable, one for signal and one
#for background
histos = {"sig" : {}, "bkg" : {}}
for var in vars :
  histos["sig"][var] = TH1F(var+"Sig",var+"Sig",50,vars[var][0],vars[var][1])
  histos["sig"][var].Sumw2() #Tell the histogram it should keep the errors on each point
  histos["sig"][var].SetMarkerColor(4) #...background is blue
  histos["sig"][var].SetLineColor(4)
  histos["sig"][var].SetMarkerStyle(20)
  histos["sig"][var].GetXaxis().SetTitle(var)
  histos["bkg"][var] = TH1F(var+"Bkg",var+"Bkg",50,vars[var][0],vars[var][1])
  histos["bkg"][var].Sumw2() #Tell the histogram it should keep the errors on each point
  histos["bkg"][var].SetMarkerColor(2) #Signal is red...
  histos["bkg"][var].SetLineColor(2)
  histos["bkg"][var].SetMarkerStyle(24)
  histos["bkg"][var].GetXaxis().SetTitle(var)


#Fill the histograms depending on whether the mass of the event
#is in the signal or the background ranges
for entry in range(0,inputchain.GetEntries()) :
  inputchain.GetEntry(entry)
  tofill = ""
  #If I am in the signal box fill the signal histograms
  if inputchain.__getattr__(nameofmassvar) > sigrangelow and \
     inputchain.__getattr__(nameofmassvar) < sigrangehigh :
    tofill = "sig"
  #Else if I am in the sideband fill the background histograms
  elif inputchain.__getattr__(nameofmassvar) > backgroundrangelow and \
       inputchain.__getattr__(nameofmassvar) < backgroundrangehigh :
    tofill = "bkg"
  #Else fill neither and move to the next event
  else : 
    continue
  for var in vars :
    histos[tofill][var].Fill(inputchain.__getattr__(var))

#Subtract the background histograms from the signal ones corrected by the scale factor
for var in vars :
  histos["sig"][var].Add(histos["bkg"][var],-1.*backscale)


#Define a canvas on which to plot all of this. What we want to do is plot the
#background subtracted signal distributions and the pure background distributions
#normalized on top of each other.
mycanvas = rootnotes.default_canvas(name="c5", size=(1024, 768))

from ROOT import TLegend
mycanvas.Divide(2,2)


currentpad = 1
for var in vars :
  mycanvas.cd(currentpad)
  #Which is taller?
  topbkg =  histos["bkg"][var].GetBinContent(histos["bkg"][var].GetMaximumBin())
  topsig =  histos["sig"][var].GetBinContent(histos["sig"][var].GetMaximumBin())
  if (topbkg > topsig) :
    histos["bkg"][var].DrawNormalized("E")
    histos["sig"][var].DrawNormalized("ESAME")
  else : 
    histos["sig"][var].DrawNormalized("E")
    histos["bkg"][var].DrawNormalized("ESAME")
  #Set to log scale
  mycanvas.GetPad(currentpad)
  legend = TLegend(0.575,0.725,0.875,0.875)
  legend.AddEntry(histos["bkg"][var], "Background", "p")
  legend.AddEntry(histos["sig"][var], "Signal", "p")
  legend.SetTextSize(0.03);
  legend.Draw()
  mycanvas.GetPad(currentpad).SetLogy()
  currentpad += 1


mycanvas

from ROOT import TF1, TLatex, kGreen, kBlack, kTRUE, Form

mycanvas.cd(1);
fFitExpo = TF1("fitExpo","expo(0)",0, 30);

fFitExpo.SetLineColor(kGreen);
fFitExpo.SetLineWidth(2);



histos["sig"]["D0_TAU"].Fit(fFitExpo,"rmeh")
labels2 = TLatex(0,0,"Fit Parameters");
labels2.SetTextSize(0.04);
labels2.SetTextColor(kBlack);
labels2.SetNDC(kTRUE);
labels2.DrawLatex(
        0.3,
        0.8,
        Form(
            "D^{0} lifetime %.4f #pm %.4f (ps)" % (
                -1./fFitExpo.GetParameter(1),
                fFitExpo.GetParError(1)/(fFitExpo.GetParameter(1) * fFitExpo.GetParameter(1))
            )
        )
    )


mycanvas.Update()

mycanvas