import ROOT

ROOT.gROOT.SetBatch(True)

# PART 1: Background-only fits
# ----------------------------

# Build plain exponential model
x = ROOT.RooRealVar("x", "x", 10, 100)
alpha = ROOT.RooRealVar("alpha", "alpha", -0.04, -0.1, -0.0)
model = ROOT.RooExponential("model", "Exponential model", x, alpha)

# Define side band regions and full range
x.setRange("LEFT", 10, 20)
x.setRange("RIGHT", 60, 100)

x.setRange("FULL", 10, 100)

data = model.generate(x, 10000)

# Construct an extended pdf, which measures the event count N **on the full range**.
# If the actual domain of x that is fitted is identical to FULL, this has no affect.
#
# If the fitted domain is a subset of `FULL`, though, the expected event count is divided by
# \f[
#   \mathrm{frac} = \frac{
#     \int_{\mathrm{Fit range}} \mathrm{model}(x)  \; \mathrm{d}x }{
#     \int_{\mathrm{Full range}} \mathrm{model}(x) \; \mathrm{d}x }.
# \f]
# `N` will therefore return the count extrapolated to the full range instead of the fit range.
#
# **Note**: When using a RooAddPdf for extending the likelihood, the same effect can be achieved with
# [RooAddPdf::fixCoefRange()](https://root.cern.ch/doc/master/classRooAddPdf.html#ab631caf4b59e4c4221f8967aecbf2a65),

N = ROOT.RooRealVar("N", "Extended term", 0, 20000)
extmodel = ROOT.RooExtendPdf("extmodel", "Extended model", model, N, "FULL")


# It can be instructive to fit the above model to either the LEFT or RIGHT
# range. `N` should approximately converge to the expected number of events in
# the full range. One may try to leave out `"FULL"` in the constructor, or the
# interpretation of `N` changes.
extmodel.fitTo(data, Range="LEFT", PrintLevel=-1)
N.Print()


# If we now do a simultaneous fit to the extended model, instead of the
# original model, the LEFT and RIGHT range will each correct their local `N`
# such that it refers to the `FULL` range.
#
# This joint fit of the extmodel will include (w.r.t. the plain model fit) a product of extended terms
# \f[
#     L_\mathrm{ext} = L
#       \cdot \mathrm{Poisson} \left( N_\mathrm{obs}^\mathrm{LEFT}  | N_\mathrm{exp} / \mathrm{frac LEFT} \right)
#       \cdot \mathrm{Poisson} \left( N_\mathrm{obs}^\mathrm{RIGHT} | N_\mathrm{exp} / \mathrm{frac RIGHT} \right)
# \f]


c = ROOT.TCanvas("c", "c", 2100, 700)
c.Divide(3)
c.cd(1)

r = model.fitTo(data, Range="LEFT,RIGHT", PrintLevel=-1, Save=True)
r.Print()

frame = x.frame()
data.plotOn(frame)
model.plotOn(frame, VisualizeError=r)
model.plotOn(frame)
model.paramOn(frame, Label="Non-extended fit")
frame.Draw()

c.cd(2)

r2 = extmodel.fitTo(data, Range="LEFT,RIGHT", PrintLevel=-1, Save=True)
r2.Print()
frame2 = x.frame()
data.plotOn(frame2)
extmodel.plotOn(frame2)
extmodel.plotOn(frame2, VisualizeError=r2)
extmodel.paramOn(frame2, Label="Extended fit", Layout=(0.4, 0.95))
frame2.Draw()

# PART 2: Extending with RooAddPdf
# --------------------------------
#
# Now we repeat the above exercise, but instead of explicitly adding an extended term to a single shape pdf (RooExponential),
# we assume that we already have an extended likelihood model in the form of a RooAddPdf constructed in the form `Nsig * sigPdf + Nbkg * bkgPdf`.
#
# We add a Gaussian to the previously defined exponential background.
# The signal shape parameters are chosen constant, since the signal region is entirely in the blinded region, i.e., the fit has no sensitivity.

c.cd(3)

Nsig = ROOT.RooRealVar("Nsig", "Number of signal events", 1000, 0, 2000)
Nbkg = ROOT.RooRealVar("Nbkg", "Number of background events", 10000, 0, 20000)

mean = ROOT.RooRealVar("mean", "Mean of signal model", 40.0)
width = ROOT.RooRealVar("width", "Width of signal model", 5.0)
sig = ROOT.RooGaussian("sig", "Signal model", x, mean, width)

modelsum = ROOT.RooAddPdf("modelsum", "NSig*signal + NBkg*background", [sig, model], [Nsig, Nbkg])

# This model will automatically insert the correction factor for the
# reinterpretation of Nsig and Nnbkg in the full ranges.
#
# When this happens, it reports this with lines like the following:
# ```
# [#1] INFO:Fitting -- RooAbsOptTestStatistic::ctor(nll_modelsum_modelsumData_LEFT) fixing interpretation of coefficients of any RooAddPdf to full domain of observables
# [#1] INFO:Fitting -- RooAbsOptTestStatistic::ctor(nll_modelsum_modelsumData_RIGHT) fixing interpretation of coefficients of any RooAddPdf to full domain of observables
# ```

r3 = modelsum.fitTo(data, Range="LEFT,RIGHT", PrintLevel=-1, Save=True)
r3.Print()

frame3 = x.frame()
data.plotOn(frame3)
modelsum.plotOn(frame3)
modelsum.plotOn(frame3, VisualizeError=r3)
modelsum.paramOn(frame3, Label="S+B fit with RooAddPdf", Layout=(0.3, 0.95))
frame3.Draw()

c.Draw()

c.SaveAs("rf204b_extendedLikelihood_rangedFit.png")