#include "TFile.h"
#include "TROOT.h"
#include "TSystem.h"
#include "RooWorkspace.h"
#include "RooAbsData.h"
#include "RooRealVar.h"
#include "RooFormulaVar.h"

#include "RooStats/ModelConfig.h"
#include "RooStats/ProfileLikelihoodCalculator.h"
#include "RooStats/LikelihoodInterval.h"
#include "RooStats/LikelihoodIntervalPlot.h"

#include <iostream>
#include <cstring>
using namespace RooFit;
using namespace RooStats;
using namespace std;

// Options for the profile likelihood scan (you can adjust these)
struct ProfileLikelihoodOptions
{
   double confLevel = 0.95;
   int nScanPoints = 100;
   bool plotAsTF1 = false;
   double poiMinPlotA = 0;   // adjust as needed
   double poiMaxPlotA = 300; // adjust as needed
   bool doHypoTest = false;
   double nullValue = 0;
};
ProfileLikelihoodOptions optPL;

// This function opens your workspace, defines totSig = nsigA + nsigB as the parameter of interest,
// and then performs and plots a profile likelihood scan on totSig.
void ProfileTotalSignal(const char *infile = "simultaneousFitWorkspace.root",
                        const char *workspaceName = "wspace",
                        const char *modelConfigName = "ModelSB",
                        const char *dataName = "combData")
{
   // -------------------------------------------------------
   // Open the ROOT file and retrieve the workspace
   // -------------------------------------------------------
   TFile *file = TFile::Open(infile);
   if (!file)
   {
      cout << "File " << infile << " not found" << endl;
      return;
   }
   RooWorkspace *w = (RooWorkspace *)file->Get(workspaceName);
   if (!w)
   {
      cout << "Workspace " << workspaceName << " not found in file " << infile << endl;
      return;
   }

   // -------------------------------------------------------
   // Retrieve (or create) the ModelConfig.
   // The ModelConfig should already include the pdf and observables.
   // -------------------------------------------------------
   ModelConfig *mc = (ModelConfig *)w->obj(modelConfigName);
   if (!mc)
   {
      cout << "ModelConfig " << modelConfigName << " not found in workspace." << endl;
      return;
   }

   // -------------------------------------------------------
   // Retrieve the combined dataset from the workspace.
   // -------------------------------------------------------
   RooAbsData *data = w->data(dataName);
   if (!data)
   {
      cout << "Data " << dataName << " not found in workspace." << endl;
      return;
   }

   // -------------------------------------------------------
   // Create the ProfileLikelihoodCalculator on the dataset and model.
   // -------------------------------------------------------
   ProfileLikelihoodCalculator pl(*data, *mc);
   pl.SetConfidenceLevel(optPL.confLevel);
   LikelihoodInterval *interval = pl.GetInterval();

   // Print the 95% (or chosen) confidence interval for totSig.
   // (The first parameter of interest is totSig.)
   RooRealVar *poi_ptr = (RooRealVar *)mc->GetParametersOfInterest()->first();
   cout << "\n>>>> " << optPL.confLevel * 100 << "% interval on " << poi_ptr->GetName()
        << " is: [" << interval->LowerLimit(*poi_ptr)
        << ", " << interval->UpperLimit(*poi_ptr) << "]\n"
        << endl;

   // -------------------------------------------------------
   // Make a plot of the profile likelihood.
   // -------------------------------------------------------
   TCanvas *c = new TCanvas();
   cout << "Drawing profile likelihood plot (use fewer scan points if this is slow)..." << endl;
   LikelihoodIntervalPlot plot(interval);
   plot.SetNPoints(optPL.nScanPoints);
   plot.SetRange(optPL.poiMinPlotA, optPL.poiMaxPlotA);
   c->Modified();
   c->Update();
   TString opt;
   if (optPL.plotAsTF1)
      opt += TString("tf1");
   plot.SetMaximum(2.3); // adjust as needed
   c->Modified();
   c->Update();
   plot.Draw(opt);

   // Optionally perform a hypothesis test if requested.
   if (optPL.doHypoTest)
   {
      RooArgSet nullparams("nullparams");
      nullparams.addClone(*poi_ptr);
      nullparams.setRealValue(poi_ptr->GetName(), optPL.nullValue);
      pl.SetNullParameters(nullparams);
      auto result = pl.GetHypoTest();
      cout << "\n>>>> Hypothesis Test Result\n";
      result->Print();
   }
}

int find_val(double num, int factor = 1e1)
{
   if (num >= 1e-1)
      return factor;
   return find_val(num * 10, factor * 10);
}

#include "RooRealVar.h"
#include "RooGaussian.h"
#include "RooExponential.h"
#include "RooAddPdf.h"
#include "RooDataSet.h"
#include "RooSimultaneous.h"
#include "RooCategory.h"
#include "RooWorkspace.h"
#include "RooFormulaVar.h"
#include "RooArgSet.h"
#include "RooStats/ModelConfig.h"
#include "RooFitResult.h"
#include "TCanvas.h"
#include "TFile.h"
#include <iostream>

using namespace RooFit;
using namespace RooStats;
using namespace std;

void rf_simultaneous_runA_runB()
{
   // ============================================================
   // 1. Define nominal (global) yields and per–run nominal sub-yields.
   // ============================================================
   // Nominal numbers for each run:
   const double nSigA_nom = 0;
   const double nBkgA_nom = 900;
   const double nSigB_nom = 0;
   const double nBkgB_nom = 1000;

   // Nominal total events in each run:
   double eventsRunA = nSigA_nom + nBkgA_nom; 
   double eventsRunB = nSigB_nom + nBkgB_nom; 

   // Nominal bgk
   double nBkgTot = nBkgA_nom + nBkgB_nom;
   double nSigTot = nSigA_nom + nSigB_nom;

   RooRealVar x("x", "Observable", 100, 150);
   x.setRange("fitRange", 100, 120); 
   x.setBins(25);

   // Combined signal and single background yields and the fraction of signal, this is what we fit:
   RooRealVar nSig("nSig", "Signal Yield", nSigTot, 0, 0.5 * nBkgTot);
   RooRealVar nBkgA("nBkgA", "Background Yield - runA", nBkgA_nom, 0.8 * nBkgA_nom, 1.2 * nBkgA_nom);
   RooRealVar nBkgB("nBkgB", "Background Yield - runB", nBkgB_nom, 0.8 * nBkgB_nom, 1.2 * nBkgB_nom);
   RooRealVar fracA("fracA", "Fraction of signal runA events", 0.5, 0., 1.);

   // Roo formula var for Signal Yield
   RooFormulaVar nSigA("nSigA", "nSig*fracA", RooArgSet(nSig, fracA));
   RooFormulaVar nSigB("nSigB", "nSig*(1-fracA)", RooArgSet(nSig, fracA));

   // ============================================================
   // 3. Build individual models for runA and runB.
   // ============================================================
   // Run A model:
   RooRealVar meanA("meanA", "Signal Mean (runA)", 125);
   RooRealVar sigmaA("sigmaA", "Signal Width (runA)", 2);
   meanA.setConstant(kTRUE);
   sigmaA.setConstant(kTRUE);
   RooGaussian gaussA("gaussA", "Signal Gaussian (runA)", x, meanA, sigmaA);

   RooRealVar lambdaA("lambdaA", "Background Slope (runA)", -0.06);
   lambdaA.setConstant(kTRUE);
   RooExponential expA("expA", "Background Exponential (runA)", x, lambdaA);

   // The yields in the extended PDF come from our reparameterized formulas:
   RooAddPdf modelA("modelA", "RunA Model", RooArgList(gaussA, expA),
                    RooArgList(nSigA, nBkgA));

   // Run B model:
   RooRealVar meanB("meanB", "Signal Mean (runB)", 125);
   RooRealVar sigmaB("sigmaB", "Signal Width (runB)", 4);
   meanB.setConstant(kTRUE);
   sigmaB.setConstant(kTRUE);
   RooGaussian gaussB("gaussB", "Signal Gaussian (runB)", x, meanB, sigmaB);

   RooRealVar lambdaB("lambdaB", "Background Slope (runB)", -0.09);
   lambdaB.setConstant(kTRUE);
   RooExponential expB("expB", "Background Exponential (runB)", x, lambdaB);

   RooAddPdf modelB("modelB", "RunB Model", RooArgList(gaussB, expB),
                    RooArgList(nSigB, nBkgB));

   // ============================================================
   // 4. Build the simultaneous model.
   RooCategory runCat("runCat", "Run Category");
   runCat.defineType("RunA");
   runCat.defineType("RunB");

   RooSimultaneous simPdf("simPdf", "Simultaneous PDF", runCat);
   simPdf.addPdf(modelA, "RunA");
   simPdf.addPdf(modelB, "RunB");

   // runCat.Print("a");
   // simPdf.Print("t");

   RooDataSet *combData0 = simPdf.generate(RooArgSet(x, runCat), Extended());
   // RooDataSet *combData = (RooDataSet *)combData0->reduce("fitRange");

   // Print the number of entries per category
   cout << "\n";
   std::cout << "Entries for RunA: " 
          << combData->sumEntries("runCat==runCat::RunA") 
          << " original number of entries: " << eventsRunA << std::endl;
std::cout << "Entries for RunB: " 
          << combData->sumEntries("runCat==runCat::RunB") 
          << " original number of entries: " << eventsRunB << std::endl;
   cout << "Entries total: " << combData->sumEntries()
      << "   original number of entries " << nBkgTot + nSigTot << endl;
   cout << "\n";

   x.setRange(120, 130); // avoid using Range("fitRange");
   RooFitResult *fitResult = simPdf.fitTo(*combData, Save() /*, Range("fitRange")*/);
   fitResult->Print("v");

   // ============================================================
   // 6. Plot results.
   // ============================================================

   TCanvas *c1 = new TCanvas("c1", "Simultaneous Fit", 1200, 600);
   c1->Divide(3, 1); // Divide into two sections for Run A and Run B

   // Frame for Run A
   x.setRange(100, 150); // reset full range for plotting
   c1->cd(1);
   RooPlot *xframeA = x.frame(Title("Run A"));
   combData->plotOn(xframeA, Cut("runCat==runCat::RunA"));
   simPdf.plotOn(xframeA, Slice(runCat, "RunA"), ProjWData(runCat, *combData), LineColor(kMagenta + 2));
   simPdf.plotOn(xframeA, Slice(runCat, "RunA"), ProjWData(runCat, *combData), Components("gaussA"), LineColor(kRed), LineStyle(kDotted));
   simPdf.plotOn(xframeA, Slice(runCat, "RunA"), ProjWData(runCat, *combData), Components("expA"), LineColor(kBlue), LineStyle(kDotted));
   xframeA->Draw();
   // Frame for Run B
   c1->cd(2);
   RooPlot *xframeB = x.frame(Title("Run B"));
   combData->plotOn(xframeB, Cut("runCat==runCat::RunB"));
   simPdf.plotOn(xframeB, Slice(runCat, "RunB"), ProjWData(runCat, *combData), LineColor(kMagenta + 2));
   simPdf.plotOn(xframeB, Slice(runCat, "RunB"), ProjWData(runCat, *combData), Components("gaussB"), LineColor(kRed), LineStyle(kDashed));
   simPdf.plotOn(xframeB, Slice(runCat, "RunB"), ProjWData(runCat, *combData), Components("expB"), LineColor(kBlue), LineStyle(kDashed));
   xframeB->Draw();
   // Frame for Full Model and Dataset (combined)
   c1->cd(3);
   RooPlot *xframeFull = x.frame(Title("Full Model"));
   combData->plotOn(xframeFull);                                                                                                              // Plot the full dataset
   simPdf.plotOn(xframeFull, ProjWData(runCat, *combData), LineColor(kMagenta + 2));                                                          // Plot the full model
   simPdf.plotOn(xframeFull, Components("gaussA"), Slice(runCat, "RunA"), ProjWData(runCat, *combData), LineColor(kRed), LineStyle(kDotted)); // Run A signal
   simPdf.plotOn(xframeFull, Components("expA"), Slice(runCat, "RunA"), ProjWData(runCat, *combData), LineColor(kBlue), LineStyle(kDotted));  // Run A background
   simPdf.plotOn(xframeFull, Components("gaussB"), Slice(runCat, "RunB"), ProjWData(runCat, *combData), LineColor(kRed), LineStyle(kDashed)); // Run B signal
   simPdf.plotOn(xframeFull, Components("expB"), Slice(runCat, "RunB"), ProjWData(runCat, *combData), LineColor(kBlue), LineStyle(kDashed));  // Run B background

   simPdf.paramOn(xframeFull);
   xframeFull->Draw();

   /*
   // ============================================================
   // 7. Save model and data.
   // ============================================================
   RooWorkspace ws("wspace", "Workspace");
   ws.import(*combData);

   ModelConfig mc("ModelSB", &ws);
   mc.SetPdf(simPdf);
   mc.SetObservables(x);
   mc.SetParametersOfInterest(RooArgSet(nSig));
   mc.SetNuisanceParameters(RooArgSet(nBkgA, nBkgB, fracA));
   mc.SetGlobalObservables(RooArgSet());
   ws.import(mc);

   ws.writeToFile("simultaneousFitWorkspace.root");

   cout << "Workspace saved to file: simultaneousFitWorkspace_modified.root" << endl;
   ws.Print("t");
   */
}