import ROOT
from ROOT import RooFit, RooRealVar, RooGaussian, RooDataSet, RooArgSet, RooArgList
from ROOT import RooAddition, RooMinimizer, RooAddPdf

def simultaneous_fit_with_roo_minimizer():
    # 1) Define a shared observable
    x = RooRealVar("x", "x", -10, 10)

    # 2) Define SHARED parameter(s), e.g. one mean for both Gaussians
    mean = RooRealVar("mean", "shared mean", 0, -10, 10)  # common
    mean2 = RooRealVar("mean2", "shared mean2", 2, -10, 10)  # common

    # Distinct widths
    sigma1 = RooRealVar("sigma1", "width for data1", 1.5, 0.1, 5.0)
    sigma2 = RooRealVar("sigma2", "width for data2", 0.5, 0.1, 4.0)
    sigma3 = RooRealVar("sigma3", "width for data3", 1.5, 0.1, 5.0)
    sigma4 = RooRealVar("sigma4", "width for data4", 2.5, 0.1, 5.0)
    sigma5 = RooRealVar("sigma5", "width for data5", 3.5, 0.1, 5.0)
    sigma6 = RooRealVar("sigma6", "width for data6", 0.5, 0.1, 5.0)

    # yields
    yield1 = RooRealVar("yield1", "yield for data1", 100000, 0, 500000)
    yield2 = RooRealVar("yield2", "yield for data2", 20000, 0, 500000)
    yield3 = RooRealVar("yield3", "yield for data3", 300000, 0, 500000)
    yield4 = RooRealVar("yield4", "yield for data4", 400000, 0, 500000)
    yield5 = RooRealVar("yield5", "yield for data5", 500000, 0, 500000)
    yield6 = RooRealVar("yield6", "yield for data6", 600000, 0, 500000)

    # 3) Build two separate PDFs that share 'mean' but differ in sigma
    _pdf1 = RooGaussian("_pdf1", "_pdf1", x, mean, sigma1)
    _pdf2 = RooGaussian("_pdf2", "_pdf2", x, mean, sigma2)
    _pdf3 = RooGaussian("_pdf3", "_pdf3", x, mean, sigma3)
    _pdf4 = RooGaussian("_pdf4", "_pdf4", x, mean2, sigma4)
    _pdf5 = RooGaussian("_pdf5", "_pdf5", x, mean, sigma5)
    _pdf6 = RooGaussian("_pdf6", "_pdf6", x, mean2, sigma6)

    pdf1 = RooAddPdf("pdf1", "pdf1", RooArgList(_pdf1, _pdf2, _pdf3, _pdf4), RooArgList(yield1, yield2, yield3, yield4))
    pdf2 = RooAddPdf("pdf2", "pdf2", RooArgList(_pdf5, _pdf6), RooArgList(yield5, yield6))

    # 4) Generate two separate data sets from these shapes
    data1 = pdf1.generate(ROOT.RooArgSet(x), 1.5 * 300000)  # "first bin"
    data2 = pdf2.generate(ROOT.RooArgSet(x), 1.5 * 400000)  # "second bin"

    # 5) Create negative log-likelihood (NLL) for each dataset
    #    Note that 'mean' is shared, so the Minimizer sees it as common
    nll1 = pdf1.createNLL(data1, RooFit.NumCPU(4), RooFit.EvalBackend("cpu")) # no multi-cores
    nll2 = pdf2.createNLL(data2, RooFit.NumCPU(4), RooFit.EvalBackend("cpu")) # no multi-cores
    # nll1 = pdf1.createNLL(data1, RooFit.NumCPU(4), RooFit.EvalBackend("legacy")) # multi-cores enabled
    # nll2 = pdf2.createNLL(data2, RooFit.NumCPU(4), RooFit.EvalBackend("legacy")) # multi-cores enabled

    ROOT.Math.MinimizerOptions.SetDefaultMinimizer("Minuit")

    # 6) Sum the two NLL objects with RooAddition to get the total NLL
    NllList = RooArgList(nll1, nll2)
    total_nll = RooAddition("total_nll", "sum of nll1 + nll2", RooArgList(NllList))

    # 7) Create a RooMinimizer on total_nll and run
    minim = RooMinimizer(total_nll)
    minim.setPrintLevel(1)
    # Minimization
    migradStatus = minim.migrad()
    hesseStatus = minim.hesse()

if __name__ == "__main__":
    simultaneous_fit_with_roo_minimizer()