import ROOT

ROOT.gROOT.SetBatch(False)

ws = ROOT.RooWorkspace("ws")
x = ws.factory("x[100, 200]")
pdf = ws.factory("RooBernstein::pdf(x, {par_0[2, -10, 10], par_1[1, -10, 10], 1})")
# no problem with exponential
#pdf = ws.factory("RooExponential::pdf(x, par_0[-0.002, -10, 10])")
# extend the pdf
model = ws.factory("SUM::model(n[40000, 0, 200000] * pdf)")
# generate dataset in the full range
data = model.generate(ROOT.RooArgSet(x))

blind_region = [140, 160]
x.setRange("left", 100, blind_region[0])
x.setRange("right", blind_region[1], 200)
# data in the sidebands
blinded_data = data.reduce("x < 140 || x > 160")

canvas = ROOT.TCanvas()
canvas.Divide(3, 1)
canvas.cd(1)

frame_1 = x.frame()
frame_1.SetTitle("generate, fit full range, plot full range")
model.fitTo(data)
data.plotOn(frame_1)
model.plotOn(frame_1)
model.paramOn(frame_1)

frame_1.Draw()

canvas.cd(2)
frame_2 = x.frame()
frame_2.SetTitle("don't fit, plot in sidebands")
blinded_data.plotOn(frame_2)
model.plotOn(frame_2, ROOT.RooFit.NormRange("left,right"), ROOT.RooFit.Range("left,right"))
model.paramOn(frame_2)

frame_2.Draw()


canvas.cd(3)
model.fitTo(blinded_data, ROOT.RooFit.Range("left,right"))
frame_3 = x.frame()
frame_3.SetTitle("fit in sidebands, plot on sidebands")
blinded_data.plotOn(frame_3)
model.plotOn(frame_3, ROOT.RooFit.NormRange("left,right"), ROOT.RooFit.Range("left,right"))
model.paramOn(frame_3)

frame_3.Draw()

canvas.Draw()

input()