import ROOT


ROOT.RooMsgService.instance().setGlobalKillBelow(ROOT.RooFit.WARNING)

# Create a histogram with 100 bins, covering the range from 0 to 10

hist = ROOT.TH1F("hist", "Histogram with Two Gaussians and Two Backgrounds", 100, 0, 10)

# Parameters for the first Gaussian
mean1 = 2
sigma1 = 0.5
amplitude1 = 50

# Proportionality factor for the second Gaussian
f = 0.6  # Adjust this factor as needed

# Parameters for the second Gaussian
mean2 = 8
sigma2 = 0.7
amplitude2 = f*amplitude1

# Parameters for the first linear background (0 to 5)
slope1 = -1
intercept1 = 5

# Parameters for the second linear background (5 to 10)
slope2 = 1
intercept2 = -5

# Fill the histogram with the Gaussians and linear backgrounds
for i in range(100):
    x = hist.GetBinCenter(i+1)  # Get the center of the bin
    if x <= 5:
        # First Gaussian
        gaussian1 = amplitude1 * ROOT.TMath.Gaus(x, mean1, sigma1)
        # First linear background
        background1 = slope1 * x + intercept1
        hist.SetBinContent(i+1, gaussian1 + background1)
    else:  # In the region 5 to 10
        # Second Gaussian
        gaussian2 = amplitude2 * ROOT.TMath.Gaus(x, mean2, sigma2)
        # Second linear background
        background2 = slope2 * x + intercept2
        hist.SetBinContent(i+1, gaussian2 + background2)

# Draw the histogram
canvas = ROOT.TCanvas('c', 'c')
hist.Draw()
canvas.Modified()
canvas.Update()

try:
    input("Press Enter to Continue")
except SyntaxError:
    pass


x = ROOT.RooRealVar('x', 'x', 0, 10)

roo_datahist = ROOT.RooDataHist("roo_datahist", "RooDataHist from TH1", ROOT.RooArgList(x), ROOT.RooFit.Import(hist))

canvas.Clear()
x_frame = x.frame()
roo_datahist.plotOn(x_frame)
x_frame.Draw()
canvas.Modified()
canvas.Update()

try:
    input("Press Enter to continue")
except SyntaxError:
    pass

x.setRange("range1", 0, 4.9)
x.setRange("range2", 5.1,10)

subhist_1 = roo_datahist.reduce(ROOT.RooFit.CutRange("range1"))
subcounts_1 = subhist_1.sumEntries()

subhist_2 = roo_datahist.reduce(ROOT.RooFit.CutRange("range2"))
subcounts_2 = subhist_2.sumEntries()

#constructing first gaussian pdf:
mean1 = ROOT.RooRealVar("mean1", "mean1", 2, 0, 5)
sigma1 = ROOT.RooRealVar("sigma1", "sigma1", 0.5, 0.1, 2)
gauss1 = ROOT.RooGaussian("gauss1", "gauss1", x, mean1, sigma1)
gauss1_yield = ROOT.RooRealVar("gauss1_yield", "gauss1_yield", 100, 0, subcounts_1)

slope1 =  ROOT.RooRealVar("slope1", "slope1", -1.0, -2, 0)
intercept1 = ROOT.RooRealVar("intercept1", "intercept1", 5, 0, 10)
bg1 = ROOT.RooPolynomial("bg1", "bg1", x, ROOT.RooArgList(intercept1,slope1), lowestOrder = 0)
bg1_yield = ROOT.RooRealVar("bg1_yield", "bg1_yield", 100, 0, subcounts_1)


pdf1 = ROOT.RooAddPdf("pdf1", "pdf1", ROOT.RooArgList(gauss1, bg1), ROOT.RooArgList(gauss1_yield, bg1_yield))

#constructing second gaussian pdf:
mean2 = ROOT.RooRealVar("mean2", "mean2", 8, 5, 10)
sigma2 = ROOT.RooRealVar("sigma2", "sigma2", 0.7, 0.1, 2)
gauss2 = ROOT.RooGaussian("gauss2", "gauss2", x, mean2, sigma2)
fraction = ROOT.RooRealVar("fraction", "fraction", 0.6, 0., 1.)
gauss2_yield = ROOT.RooFormulaVar("gauss2_yield", "@0*@1", ROOT.RooArgList(fraction, gauss1_yield))

slope2 =  ROOT.RooRealVar("slope2", "slope2", 1.0, 0, 2)
intercept2 = ROOT.RooRealVar("intercept2", "intercept2", -5, -10, 0)
bg2 = ROOT.RooPolynomial("bg2", "bg2", x, ROOT.RooArgList(intercept2,slope2), lowestOrder = 0)
bg2_yield = ROOT.RooRealVar("bg2_yield", "bg2_yield", 100, 0, subcounts_2)


pdf2 = ROOT.RooAddPdf("pdf2", "pdf2", ROOT.RooArgList(gauss2, bg2), ROOT.RooArgList(gauss2_yield, bg2_yield))

# Define category to distinguish left and right events
sample = ROOT.RooCategory("sample", "sample")
sample.defineType("left")
sample.defineType("right")
 
# Construct combined dataset in (x,sample)
combData = ROOT.RooDataSet(
    "combData",
    "combined data",
    {x},
    Index=sample,
    Import={"left": subhist_1, "right": subhist_2},
)

simPdf = ROOT.RooSimultaneous("simPdf", "simultaneous pdf", {"left": pdf1, "right": pdf2}, sample)
 
fit_result = simPdf.fitTo(combData, PrintLevel=-1, Save=True)
fit_result.Print()


x_frame = x.frame()
simPdf.plotOn(x_frame, ProjWData=(sample, combData))
combData.plotOn(x_frame)
x_frame.Draw()
canvas.Modified()
canvas.Update()

try:
    input("Press Enter to continue")
except SyntaxError:
    pass