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

inputfile = TFile("TESHEP_Pedagogical_Data.root")

inputchain = inputfile.Get("DecayTree")


nameoffitvar = "D0_MM"


rangeoffit = (1816., 1914)


backgroundrange = (1890., 1914.)


nsigmasigbox = 2.


massvar = RooRealVar("D0_MM", "D0_MM", rangeoffit[0], rangeoffit[1])


mydataset = RooDataSet("DecayTreeDataSet","DecayTreeDataSet",RooArgSet(massvar))


for entry in range(0,inputchain.GetEntries()) :
  inputchain.GetEntry(entry)
  if inputchain.__getattr__(nameoffitvar) > rangeoffit[1] or \
     inputchain.__getattr__(nameoffitvar) < rangeoffit[0] :
    continue
  massvar.setVal(inputchain.__getattr__(nameoffitvar))
  mydataset.add(RooArgSet(massvar))


gausMean_B 	= RooRealVar("gausMean_B","gausMean_B",(rangeoffit[0]+rangeoffit[1])/2.,rangeoffit[0],rangeoffit[1])
gausWidth_B 	= RooRealVar("gausWidth_B","gausWidth_B",15.,1.,50.)
chebOne_B 	= RooRealVar("chebOne_B","chebOne_B",0.1,-5.,5.)
gaussian_B 	= RooGaussian("gaussian_B","gaussian_B",massvar,gausMean_B,gausWidth_B)
cheb_B     	= RooChebychev("cheb_B","cheb_B",massvar,RooArgList(chebOne_B))

N_Sig_B = RooRealVar("N_Sig_B","N_Sig_B",10000,0,250000)
N_Bkg_B = RooRealVar("N_Bkg_B","N_Bkg_B",10000,0,500000)


totalpdf_B = RooAddPdf("totalpdf_B","totalpdf_B", RooArgList(gaussian_B,cheb_B), RooArgList(N_Sig_B,N_Bkg_B))
totalpdf_B.fitTo(mydataset, RooFit.Extended())


canv_Dmass = rootnotes.default_canvas()
frame_Dmass = massvar.frame()
mydataset.plotOn(frame_Dmass,RooFit.Binning(100))
totalpdf_B.plotOn(frame_Dmass)
frame_Dmass.Draw()

canv_Dmass

massvar.setRange("integralRangeAll", rangeoffit[0], rangeoffit[1])
massvar.setRange("integralRangeSignal", 
                 gausMean_B.getVal() - nsigmasigbox * gausWidth_B.getVal(),
                 gausMean_B.getVal() + nsigmasigbox * gausWidth_B.getVal())
massvar.setRange("integralRangeBackground", backgroundrange[0], backgroundrange[1])


integralall    = cheb_B.createIntegral(RooArgSet(massvar), RooFit.NormSet(RooArgSet(massvar)), RooFit.Range("integralRangeAll"))
integralsignal = cheb_B.createIntegral(RooArgSet(massvar), RooFit.NormSet(RooArgSet(massvar)), RooFit.Range("integralRangeSignal"))
integralback   = cheb_B.createIntegral(RooArgSet(massvar), RooFit.NormSet(RooArgSet(massvar)), RooFit.Range("integralRangeBackground"))

normalizedIntegralValue_S = integralsignal.getVal()/integralall.getVal()
normalizedIntegralValue_B = integralback.getVal()/integralall.getVal()


print "Signal range =",\
    gausMean_B.getVal() - nsigmasigbox * gausWidth_B.getVal(), "-",\
    gausMean_B.getVal() + nsigmasigbox*gausWidth_B.getVal()
print "Background in signal range =", normalizedIntegralValue_S * N_Bkg_B.getVal()
print "Background in background range =",normalizedIntegralValue_B * N_Bkg_B.getVal()

print "Background scale factor =", normalizedIntegralValue_S / normalizedIntegralValue_B