#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
An example of a simple 1D fit using RooFit in python
"""
__author__ = "Sanjeeda"

import ROOT as r

# turns off popup plots
r.gROOT.SetBatch(True)

# get the data in a TTree format
fi = r.TFile('forfit_proc11_bucket9to12_bcs.root', 'READ')
tr = r.TTree()
fi.GetObject('after_bcs', tr)

# check we have the tree
print(tr.GetEntries())

# get the data in RooFit land
deltaM = r.RooRealVar('deltaM_after_bcs', 'deltaM', 0.139, 0.159, 'GeV/c^{2}')
rds = r.RooDataSet(
    'data', 'D^{0}#rightarrow K^{0}_{S} K^{0}_{S}', tr, r.RooArgSet(deltaM))

# always draw plot to make sure things look sane -- saves a lot of time debugging
canvas = r.TCanvas()
plot = deltaM.frame()
plot.SetTitle('Plot of the data only')
rds.plotOn(plot)
plot.Draw()
canvas.SaveAs('data-only.pdf')

#signal model
mean = r.RooRealVar('mean', 'mean_of_the_dist', 0.145, 0.139, 0.159, 'GeV/c^{2}')
sigma = r.RooRealVar('sigma', 'sigma_of_the_dist', 0.006, 0.0, 1)
sig = r.RooGaussian('sig', 'signal component',deltaM, mean, sigma)

# background model
#alpha = r.RooRealVar('alpha', 'co_efficient',20, 'GeV/c^{2}')
#bkg = r.RooGenericPdf("bkg","pow((deltaM -0.1399),1/2) + alpha*(pow((deltaM-0.1399),3/2))",r.RooArgList(deltaM, alpha))
#bkg = r.RooPolynomial("bkg","pow((deltaM -0.1399),1/2) + alpha*(pow((deltaM-0.1399),3/2))",r.RooArgList(deltaM, alpha))


# if you want to do an extended maximum likelihood fit uncomment these lines:
nsg = r.RooRealVar('nsg', 'n_{s}', 7000, 0, 10000)
nbg = r.RooRealVar('nbg', 'n_{b}', 3000, 0, 10000)
pdf = r.RooAddPdf('pdf', 'two component model', r.RooArgList(sig, bkg), r.RooArgList(nsg, nbg))

# a pre-fitting plot
debug = deltaM.frame()
debug.SetTitle('A plot without fit for debugging')
rds.plotOn(debug)
bkg.plotOn(debug)
sig.plotOn(debug)
pdf.plotOn(debug, r.RooFit.LineColor(r.kRed))
debug.Draw()
canvas.SaveAs('pre-fit.pdf')

# Run the fit.
# ------------
# The call to pdf.fitTo takes a number of "RooCmdArgs" or RooFit commands. 
#   -- "Save" tells the fitTo method to actually return the fit result in a 
#      "RooFitResult" object. Useful for accessing the fitted parameter values
#      and fit quality later.
#   -- "Strategy": by default a fast MINUT strategy is used for speed. A better
#      choice for a reliably correct minimum is strategy 2.
#      For more information see
#         https://root.cern.ch/root/htmldoc/guides/minuit2/Minuit2.html#m-strategy
#   -- "Extended": if you want to do an extended maximum likelihood fit, you
#      need to add `r.RooFit.Extended(True)` too.
rfr = pdf.fitTo(rds, r.RooFit.Save(True), r.RooFit.Strategy(2), r.RooFit.Extended(True))

# plot fitted otput
plot.SetTitle('Fitted distribution')
pdf.plotOn(plot, r.RooFit.LineColor(r.kRed))
pdf.plotOn(plot, r.RooFit.Components('sig'), r.RooFit.LineColor(r.kGreen))
pdf.plotOn(plot, r.RooFit.Components('bkg'), r.RooFit.LineColor(r.kBlue))
plot.Draw()
canvas.SaveAs('post-fit.pdf')

# figure out paramters access the value in code
rfr.Print()
pars = rfr.floatParsFinal()
pars.Print()
print(pars[0].getVal(), '+/-', pars[0].getError())