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

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

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', 'delta_M', 0.14, 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.1454, 0.14, 0.159, 'GeV/c^{2}')
sigma = r.RooRealVar('sigma', 'sigma_of_the_dist', 0.0003, 0, 10)
sig = r.RooGaussian('sig', 'signal component',deltaM, mean, sigma)

#background model
alpha = r.RooRealVar('alpha', 'co_efficient',20,-100,100)
diff = r.RooFormulaVar("diff","deltaM_after_bcs-mpi", r.RooArgList(deltaM,mpi))
first_part = r.RooFormulaVar("first_part", "pow(diff,1/2)", r.RooArgList(diff))
second_part = r.RooFormulaVar("second_part", "alpha*pow(diff,3/2)", r.RooArgList(alpha, diff))
bkg = r.RooGenericPdf("bkg","first_part + second_part",r.RooArgList(first_part, second_part))


# if you want to do an extended maximum likelihood fit uncomment these lines:
n_sig = r.RooRealVar('n_sig', 'n_{s}', 2000, 0, 10000)
n_bkg = r.RooRealVar('n_bkg', 'n_{b}', 100, 0, 10000)

pdf = r.RooAddPdf('pdf', 'two component model', r.RooArgList(sig, bkg), r.RooArgList(n_sig, n_bkg))

# 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())