#include "RooRealVar.h"
#include "RooDataSet.h"
#include "RooGaussian.h"
#include "RooBernstein.h"
#include "RooFormulaVar.h"
#include "RooAddPdf.h"
#include "RooProdPdf.h"
#include "RooCategory.h"
#include "RooSimultaneous.h"
#include "RooFitResult.h"
#include "RooPlot.h"

#include "TCanvas.h"
#include "TFile.h"

using namespace RooFit;

void fitComplete() {

    // Suppress RooFit messages
    // RooMsgService::instance().setGlobalKillBelow(RooFit::PROGRESS);

    // Output file
    TFile outFile("fitComplete.root", "RECREATE");

    // Observable
    RooRealVar x("x", "x observable", 0, 10);

    // Set global fit range for 
    // x.setRange("fitRangeGlobal", 1.0, 9.0);

    // Bin the data 
    x.setBins(50);

    // Define number of bins in each other dimensions ()
    int nBinsT = 4; 
    int nBinsE = 4;

    // Create category for slices
    RooCategory slice("slice", "slice");
    for (int t = 0; t < nBinsT; ++t) {
        for (int e = 0; e < nBinsE; ++e) {
            TString sliceName = Form("t%d_e%d", t, e);
            slice.defineType(sliceName.Data());
        }
    }

    // Create combined dataset and simultaneous PDF
    RooSimultaneous combPdf("combPdf", "combPdf", slice);
   
    // Create map for each slice data
    std::map<std::string, RooDataSet*> mapData;
    std::map<std::string, RooDataSet*> mapDataFitRange;
    
    // Create map for each slice model components 
    struct modelComponents {
        RooAbsPdf* sig;
        RooAbsPdf* bkgA;
        RooAbsPdf* bkgB;
        RooRealVar* nSig;
        RooRealVar* nBkgA;
        RooRealVar* nBkgB;
    };
    std::map<std::string, modelComponents> mapModel;

    for (int t = 0; t < nBinsT; ++t) {
        for (int e = 0; e < nBinsE; ++e) {
            TString sliceName = Form("t%d_e%d", t, e);

            // Parameters for this slice
            int nDataPoints = 10000;
            int nSignalPoints = (t - e) * 500 + 2000;
            int nBackgroundPoints = nDataPoints - nSignalPoints;
            double meanVal = 5 + ((t - e) / 1.5);

            // Set fit range for x
            TString rangeName = TString::Format("fitRange_t%d_e%d", t, e);
            double rangeMin = std::max(0.0, meanVal - 3.0);
            double rangeMax = std::min(10.0, meanVal + 3.0);
            x.setRange(rangeName, rangeMin, rangeMax);

            // Gaussian PDF in x 
            RooRealVar* mean = new RooRealVar(Form("mean_t%d_e%d", t, e), "mean", meanVal, 0, 10);
            RooRealVar* sigma = new RooRealVar(Form("sigma_t%d_e%d", t, e), "sigma", 0.3, 0.1, 5);
            RooGaussian* gaussX = new RooGaussian(Form("gauss_t%d_e%d", t, e), "Gaussian PDF", x, *mean, *sigma);

            // Bernstein polynomial background in x
            RooRealVar* a0 = new RooRealVar(Form("a0_t%d_e%d", t, e), "a0 parameter", 0.5, 0, 1);
            RooRealVar* a1 = new RooRealVar(Form("a1_t%d_e%d", t, e), "a1 parameter", 0.3, 0, 1);
            RooBernstein* polA = new RooBernstein(Form("bkgA_t%d_e%d", t, e), "Background PDF", x, RooArgList(*a0, *a1));

            // Bernstein polynomial background in x (same as above, used to show constraints)
            RooRealVar* b0 = new RooRealVar(Form("b0_t%d_e%d", t, e), "b0 parameter", 0.5, 0, 1);
            RooRealVar* b1 = new RooRealVar(Form("b1_t%d_e%d", t, e), "b1 parameter", 0.3, 0, 1);
            RooBernstein* polB = new RooBernstein(Form("bkgB_t%d_e%d", t, e), "Background PDF", x, RooArgList(*b0, *b1));

            // Create model
            RooRealVar* nSig = new RooRealVar(Form("nSig_t%d_e%d", t, e), "number of signal events", nSignalPoints, 0, nDataPoints);
            RooRealVar* nBkgA = new RooRealVar(Form("nBkgA_t%d_e%d", t, e), "number of background events", nBackgroundPoints / 2, 0, nDataPoints);
            RooRealVar* nBkgB = new RooRealVar(Form("nBkgB_t%d_e%d", t, e), "number of background events", nBackgroundPoints / 2, 0, nDataPoints);
            RooAddPdf* modelX = new RooAddPdf(Form("modelX_t%d_e%d", t, e), "Composite PDF", RooArgList(*gaussX, *polA, *polB), RooArgList(*nSig, *nBkgA, *nBkgB));

            // Add data to map
            RooDataSet* dataSlice = modelX->generate(x, nDataPoints);
            mapData[sliceName.Data()] = dataSlice;

            // Create reduced dataset in fit range
            RooDataSet* dataSliceFitRange = (RooDataSet*)dataSlice->reduce(CutRange(rangeName));
            mapDataFitRange[sliceName.Data()] = dataSliceFitRange;

            // Add components to map
            modelComponents components = {gaussX, polA, polB, nSig, nBkgA, nBkgB};
            mapModel[sliceName.Data()] = components;

            // Add PDF to simultaneous PDF
            combPdf.addPdf(*modelX, sliceName);
        }
    }

    RooDataSet combDataSet("combDataSet", "combDataSet", RooArgSet(x, slice), Index(slice), Import(mapData));
    // RooDataSet combDataSet("combDataSet", "combDataSet", RooArgSet(x, slice), Index(slice), Import(mapDataFitRange));

    // Create constraints between parameters of different slices
    RooArgList constraintList;

    for (int t = 0; t < nBinsT; ++t) {
        for (int e = 0; e < nBinsE; ++e) {
            TString sliceName = Form("t%d_e%d", t, e);

            RooRealVar* nBkgA = mapModel[sliceName.Data()].nBkgA;
            RooRealVar* nBkgB = mapModel[sliceName.Data()].nBkgB;

            double expectedFraction = 0.5 + 0.1 * (t - e); // Example expected fraction

            TString formulaName = Form("ratio_%s", sliceName.Data());
            RooFormulaVar* ratio = new RooFormulaVar(formulaName, TString::Format("@0/(@0+@1)-%f", expectedFraction), RooArgList(*nBkgA, *nBkgB));
            
            RooRealVar* sigma = new RooRealVar(Form("sigma_constraint_%s", sliceName.Data()), "constraint sigma", 0.05, 0.001, 1);
            RooRealVar* mean = new RooRealVar(Form("mean_constraint_%s", sliceName.Data()), "constraint mean", 0.0, -1, 1);
            sigma->setConstant(kTRUE);
            mean->setConstant(kTRUE);

            TString constraintName = Form("constraint_%s", sliceName.Data());
            RooGaussian* constraint = new RooGaussian(constraintName, constraintName, *ratio, *mean, *sigma);

            constraintList.add(*constraint);
        }
    }

    // // Fit the model with constraints:
    RooFitResult* fitResult = combPdf.fitTo(combDataSet, Save(), PrintLevel(-1), Extended(), Minos(false), ExternalConstraints(constraintList));

    // Fit the model with constraints and global fit range:
    // RooFitResult* fitResult = combPdf.fitTo(combDataSet, Save(), PrintLevel(-1), Extended(), Minos(false), ExternalConstraints(constraintList), Range("fitRangeGlobal"));

    // Fit the model with constraints and individual fit ranges in each slice:
    // TString rangeString;
    // for (int t = 0; t < nBinsT; ++t) {
    //     for (int e = 0; e < nBinsE; ++e) {
    //         TString rangeName = Form("fitRange_t%d_e%d", t, e);
    //         if (t == 0 && e == 0) {
    //             rangeString = rangeName;
    //         } else {
    //             rangeString += "," + rangeName;
    //         }
    //     }
    // }
    // RooFitResult* fitResult = combPdf.fitTo(combDataSet, Save(), PrintLevel(-1), Extended(), Minos(false), ExternalConstraints(constraintList), Range(rangeString)); // SplitRange(true)?
   
    fitResult->Print("v");

    // Create canvases for prefit and fit plots
    TCanvas *cFitAll = new TCanvas("cFitAll", "cFitAll", 1200, 800);
    cFitAll->Divide(nBinsE, nBinsT);

    for (int t = 0; t < nBinsT; ++t) {
        for (int e = 0; e < nBinsE; ++e) {
            TString sliceName = Form("t%d_e%d", t, e);

            RooPlot* frame = x.frame(RooFit::Title(sliceName));

            // Plot data for this slice
            combDataSet.plotOn(frame, Cut(Form("slice==slice::%s", sliceName.Data())));
            
            // Draw without fit range:
            combPdf.plotOn(frame, Slice(slice, sliceName), ProjWData(slice, combDataSet), LineColor(kBlue));
            combPdf.plotOn(frame, Slice(slice, sliceName), ProjWData(slice, combDataSet), Components(*mapModel[sliceName.Data()].sig), LineColor(kGreen+2));
            combPdf.plotOn(frame, Slice(slice, sliceName), ProjWData(slice, combDataSet), Components(*mapModel[sliceName.Data()].bkgA), LineColor(kRed), LineStyle(kDashed));
            combPdf.plotOn(frame, Slice(slice, sliceName), ProjWData(slice, combDataSet), Components(*mapModel[sliceName.Data()].bkgB), LineColor(kOrange), LineStyle(kDotted));

            // Draw with global fit range:
            // combPdf.plotOn(frame, Slice(slice, sliceName), ProjWData(slice, combDataSet), Range("fitRangeGlobal"), NormRange("fitRangeGlobal"), LineColor(kBlue));
            // combPdf.plotOn(frame, Slice(slice, sliceName), ProjWData(slice, combDataSet), Components(*mapModel[sliceName.Data()].sig), Range("fitRangeGlobal"), NormRange("fitRangeGlobal"), LineColor(kGreen+2));
            // combPdf.plotOn(frame, Slice(slice, sliceName), ProjWData(slice, combDataSet), Components(*mapModel[sliceName.Data()].bkgA), Range("fitRangeGlobal"), NormRange("fitRangeGlobal"), LineColor(kRed), LineStyle(kDashed));
            // combPdf.plotOn(frame, Slice(slice, sliceName), ProjWData(slice, combDataSet), Components(*mapModel[sliceName.Data()].bkgB), Range("fitRangeGlobal"), NormRange("fitRangeGlobal"), LineColor(kOrange), LineStyle(kDotted));

            // Try method from: https://root-forum.cern.ch/t/simultaneous-fit-normalization-issue/33965/2
            // const double entries = combDataSet.sumEntries(Form("slice==slice::%s", sliceName.Data()));
            // const double entries = combDataSet.sumEntries();
            // combPdf.plotOn(frame, Slice(slice, sliceName), ProjWData(slice, combDataSet), LineColor(kBlue), Range("fitRangeGlobal"), Normalization(entries, RooAbsReal::NumEvent));
            // combPdf.plotOn(frame, Slice(slice, sliceName), ProjWData(slice, combDataSet), Components(*mapModel[sliceName.Data()].sig), Range("fitRangeGlobal"), LineColor(kGreen+2), Normalization(entries, RooAbsReal::NumEvent));
            // combPdf.plotOn(frame, Slice(slice, sliceName), ProjWData(slice, combDataSet), Components(*mapModel[sliceName.Data()].bkgA), Range("fitRangeGlobal"), LineColor(kRed), LineStyle(kDashed), Normalization(entries, RooAbsReal::NumEvent));
            // combPdf.plotOn(frame, Slice(slice, sliceName), ProjWData(slice, combDataSet), Components(*mapModel[sliceName.Data()].bkgB), Range("fitRangeGlobal"), LineColor(kOrange), LineStyle(kDotted), Normalization(entries, RooAbsReal::NumEvent));

            // Draw with individual fit ranges:
            // TString rangeName = Form("fitRange_t%d_e%d", t, e);
            // combPdf.plotOn(frame, Slice(slice, sliceName), ProjWData(slice, combDataSet), Range(rangeName), NormRange(rangeName), LineColor(kBlue));
            // combPdf.plotOn(frame, Slice(slice, sliceName), ProjWData(slice, combDataSet), Components(*mapModel[sliceName.Data()].sig), Range(rangeName), NormRange(rangeName), LineColor(kGreen+2));
            // combPdf.plotOn(frame, Slice(slice, sliceName), ProjWData(slice, combDataSet), Components(*mapModel[sliceName.Data()].bkgA), Range(rangeName), NormRange(rangeName), LineColor(kRed), LineStyle(kDashed));
            // combPdf.plotOn(frame, Slice(slice, sliceName), ProjWData(slice, combDataSet), Components(*mapModel[sliceName.Data()].bkgB), Range(rangeName), NormRange(rangeName), LineColor(kOrange), LineStyle(kDotted));

            int pad = (nBinsT - 1 - t) * nBinsE + e + 1;
            cFitAll->cd(pad);
            frame->Draw();

            TCanvas *cFit = new TCanvas(Form("cFit_%s", sliceName.Data()), Form("Fit for %s", sliceName.Data()), 600, 400);
            frame->Draw();
            cFit->Write();
        }
    }
    
    // Write canvases to output file
    cFitAll->Write();

    // Close output file
    outFile.Close();
}