#include <iostream>
#include <vector>
#include <string>
#include "TFile.h"
#include "TTree.h"
#include "RooRealVar.h"
#include "RooArgSet.h"
#include "RooArgList.h"
#include "RooDataSet.h"
#include "RooAbsPdf.h"
#include "RooProdPdf.h"
#include "RooAddPdf.h"
#include "RooWorkspace.h"
#include "RooFitResult.h"
#include "RooCmdArg.h"
#include "RooConstVar.h"

using namespace std;
using namespace RooFit;

void createwscommon()
{

    // Define years to process
    vector<int> year = {2021, 2022};
    // Names used in the workspaces and PDFs
    string ws_name = "ws";
    string model_name = "model";
    string ric_name = "ric_rmdacc";
    string time_name = "time";
    string minv_name = "minv";

    // Loop over years
    for (int i = 0; i < year.size(); i++)
    {
        string year_string = to_string(year[i]);
        cout << "\nProcessing year " << year_string << "\n"
             << endl;

        // File names for the workspaces
        string ws_rmdfitname = "ws_rmdfit_" + year_string + ".root";
        string ws_accfitname = "ws_accfit_" + year_string + ".root";
        string ws_gausflattimename = "ws_gausflattime_" + year_string + ".root";
        string ws_polytimename = "ws_polytime_" + year_string + ".root";

        // Open the workspace files
        TFile *file_rmd = new TFile(ws_rmdfitname.c_str());
        TFile *file_acc = new TFile(ws_accfitname.c_str());
        TFile *file_gausflattime = new TFile(ws_gausflattimename.c_str());
        TFile *file_polytime = new TFile(ws_polytimename.c_str());

        // Get the workspaces
        RooWorkspace *ws_rmd = (RooWorkspace *)file_rmd->Get(ws_name.c_str());
        RooWorkspace *ws_acc = (RooWorkspace *)file_acc->Get(ws_name.c_str());
        RooWorkspace *ws_gausflattime = (RooWorkspace *)file_gausflattime->Get(ws_name.c_str());
        RooWorkspace *ws_polytime = (RooWorkspace *)file_polytime->Get(ws_name.c_str());

        // For debugging: print the workspaces
        ws_rmd->Print("v");
        ws_acc->Print("v");
        ws_gausflattime->Print("v");
        ws_polytime->Print("v");

        // Retrieve the PDFs from the workspaces.
        // We assume that ws_rmd and ws_acc contain PDFs named "model_rmd" and "model_acc"
        // and that ws_gausflattime contains PDFs "gaus" and "cte".
        RooAbsPdf *pdf_rmd = (RooAbsPdf *)ws_rmd->pdf((model_name + "_rmd").c_str());
        RooAbsPdf *pdf_acc = (RooAbsPdf *)ws_acc->pdf((model_name + "_acc").c_str());
        RooAbsPdf *pdf_gaus = (RooAbsPdf *)ws_gausflattime->pdf("gaus_time");
        RooAbsPdf *pdf_cte = (RooAbsPdf *)ws_gausflattime->pdf("cte_time");
        RooAbsPdf *pdf_poly = (RooAbsPdf *)ws_polytime->pdf("poly_time");

        // Create a new (common) workspace to host the shared observables.
        // (The idea is to later update the observables inside the model.)
        TFile *outws = new TFile(("ws_common_" + year_string + ".root").c_str(), "RECREATE");
        RooWorkspace *ws_common = new RooWorkspace("ws", "ws");

        // Import the cloned PDFs into the common workspace.
        ws_common->import(*pdf_rmd);
        ws_common->import(*pdf_acc);
        ws_common->import(*pdf_gaus);
        ws_common->import(*pdf_cte);
        ws_common->import(*pdf_poly);

        ws_common->Print("v");

        // Write the common workspace to the output file and close the file
        ws_common->Write();
        outws->Close();
    }
}

RooHist *RemovePointsInRange(RooHist *originalHist, double xMin, double xMax)
{
    // Create a new RooHist to store the filtered points
    RooHist *filteredHist = new RooHist(*originalHist);

    // Loop over the points in reverse order to avoid index shifting issues
    for (int i = originalHist->GetN() - 1; i >= 0; i--)
    {
        double x, y;
        originalHist->GetPoint(i, x, y); // Get the point coordinates

        // Check if the x-value is within the range to exclude
        if (x >= xMin && x <= xMax)
        {
            // Remove the point by setting its weight to 0 (effectively "removes" it from the plot)
            filteredHist->RemovePoint(i);
        }
    }

    // Return the filtered RooHist
    return filteredHist;
}

vector<double> calculateChi2NDOF(RooHist *hist, int params)
{
    double chi2 = 0;
    int ndof = 0;

    if (ndof - params > 0)
        ndof = ndof - params;
    else
        ndof = 0;
    vector<double> chindof;

    for (int i = 0; i < hist->GetN(); i++)
    {
        double x, y;
        hist->GetPoint(i, x, y);
        double error = hist->GetErrorY(i);
        double model = 0; // Assume the model value is 0 for simplicity
        chi2 += pow(y - model, 2) / pow(error, 2);
        ndof++;
    }

    chindof.push_back(chi2);
    chindof.push_back(ndof);

    return chindof;
}

void fit2d_background()
{

    createwscommon();

    // Define years to process
    vector<int> year = {2021, 2022};
    // Names used in the workspaces and PDFs
    string ws_name = "ws";
    string model_name = "model";
    string ric_name = "ric_rmdacc";
    string time_name = "time";
    string minv_name = "minv";

    // Loop over years
    for (int i = 0; i < year.size(); i++)
    {

        // open ws_common
        string year_string = to_string(year[i]);
        string ws_common_name = "ws_common_" + year_string + ".root";
        TFile *file_common = new TFile(ws_common_name.c_str());
        RooWorkspace *ws_common = (RooWorkspace *)file_common->Get(ws_name.c_str());
        ws_common->Print("v");

        // Open the dataset file and get the tree
        string datarecfile = "rec_" + year_string + "_side_fit.root";
        TFile *dataFile = new TFile(datarecfile.c_str());
        TTree *t = (TTree *)dataFile->Get(ric_name.c_str());

        // Get the observables from the workspace
        RooRealVar *time = (RooRealVar *)ws_common->var(time_name.c_str());
        time->setBins(50);
        RooRealVar *minv = (RooRealVar *)ws_common->var(minv_name.c_str());
        minv->setBins(50);
        // Create the ranges of integration for time and minv
        double minv_min = -300, minv_max = 300;
        RooRealVar *mean_var = (RooRealVar *)ws_common->var("mean_time");
        RooRealVar *sigma_var = (RooRealVar *)ws_common->var("sigma_time");
        double mean = mean_var->getVal();
        double sigma = sigma_var->getVal();
        double time_min = mean - 3 * sigma, time_max = mean + 3 * sigma;
        time->setRange(time_min, time_max);
        time->setRange("fullRange", time_min, time_max);
        time->setRange("leftSide", time_min, time_max);
        time->setRange("rightSide", time_min, time_max);
        minv->setRange(minv_min, minv_max);
        minv->setRange("fullRange", minv_min, minv_max);
        minv->setRange("leftSide", minv_min, -84);
        minv->setRange("rightSide", 121, minv_max);

        // Create the background model
        RooAbsPdf *pdf_rmd = ws_common->pdf((model_name + "_rmd").c_str());
        RooAbsPdf *pdf_acc = ws_common->pdf((model_name + "_acc").c_str());
        RooAbsPdf *pdf_gaus = ws_common->pdf("gaus_time");
        RooAbsPdf *pdf_cte = ws_common->pdf("cte_time"); // change name from cte to poly as I want to do poly or cte

        // Create the 2d dataset
        RooDataSet data("data", "data", RooArgSet(*minv, *time), Import(*t));
        RooAbsData *dataLeft = data.reduce(CutRange("leftSide"));
        RooAbsData *dataRight = data.reduce(CutRange("rightSide"));

        // Set all variables constant
        RooArgSet *params = pdf_rmd->getParameters(RooArgSet(*minv));
        params->setAttribAll("Constant", kTRUE);
        params = pdf_acc->getParameters(RooArgSet(*minv));
        params->setAttribAll("Constant", kTRUE);
        params = pdf_gaus->getParameters(RooArgSet(*time));
        params->setAttribAll("Constant", kTRUE);
        params = pdf_cte->getParameters(RooArgSet(*time));
        params->setAttribAll("Constant", kTRUE);

        // Create the 2d acc and rmd model
        RooProdPdf model_acc((string(pdf_acc->GetName()) + "_2d").c_str(), (string(pdf_acc->GetName()) + "_2d").c_str(), RooArgList(*pdf_acc, *pdf_cte));
        RooProdPdf model_rmd((string(pdf_rmd->GetName()) + "_2d").c_str(), (string(pdf_rmd->GetName()) + "_2d").c_str(), RooArgList(*pdf_rmd, *pdf_gaus));

        // Create the variables for rmd and acc
        RooRealVar nRmd("nRmd", "nRmd", 1000, 0, 5000);
        RooRealVar nAcc("nAcc", "nAcc", 500, 0, 5000);

        // Create the 2d pdf
        RooAddPdf model((model_name + "_2d").c_str(), (model_name + "_2d").c_str(), RooArgList(model_rmd, model_acc), RooArgList(nRmd, nAcc));
        model.fixCoefNormalization(RooArgSet(*minv, *time));
        model.Print("v");

        // Fit the model to the data
        RooFitResult *result = model.fitTo(data, Save(), Range("leftSide,rightSide"), Strategy(2));
        model.removeStringAttribute("fitrange");
        result->Print("v");

        gStyle->SetOptStat(0);
        // draw dataset
        TCanvas *c3 = new TCanvas();
        c3->Divide(2, 2);
        c3->cd(2);
        TH1 *hh_data = data.createHistogram("minv,time", Binning(50), Binning(50));
        hh_data->Draw("lego");
        TH1 *hh_pdf = model.createHistogram("minv,time", 50, 50);
        hh_pdf->SetLineColor(kBlue);
        hh_pdf->SetLineWidth(2);
        hh_pdf->SetContour(30);
        hh_pdf->Draw("surf2 same");
        c3->cd(4);
        RooPlot *plot24 = minv->frame(Title(("m^{2}_{inv}    " + year_string).c_str()));
        data.plotOn(plot24);
        model.plotOn(plot24, LineColor(kMagenta + 2), Range("leftSide"), NormRange("leftSide,rightSide"));
        model.plotOn(plot24, LineColor(kBlue), Range("fullRange"), Components((model_name + "_acc").c_str()), NormRange("leftSide,rightSide"));
        model.plotOn(plot24, LineColor(kRed), Range("fullRange"), Components((model_name + "_rmd").c_str()), NormRange("leftSide,rightSide"));
        model.plotOn(plot24, LineColor(kMagenta + 2), Range("rightSide"), NormRange("leftSide,rightSide"));
        model.plotOn(plot24, LineColor(kMagenta + 2), Range("fullRange"), NormRange("leftSide,rightSide"), LineStyle(kDashed));
        model.paramOn(plot24);
        plot24->Draw();
        c3->cd(1);
        RooPlot *plot25 = time->frame(Title(("t_{e#gamma}   " + year_string).c_str()));
        data.plotOn(plot25);
        model.plotOn(plot25, LineColor(kMagenta + 2), Range("fullRange"), NormRange("fullRange"));
        model.plotOn(plot25, LineColor(kRed), Components(pdf_gaus->GetName()));
        model.plotOn(plot25, LineColor(kBlue), Components(pdf_cte->GetName()));
        data.plotOn(plot25);
        plot25->Draw();

        // Save the global time range (assumes time->getMin() and time->getMax() return the full range)
        double timeGlobalMin = time->getMin();
        double timeGlobalMax = time->getMax();
        int nslices = 2;
        double dt = (timeGlobalMax - timeGlobalMin) / double(nslices);

        // Create three ranges for the slices
        for (int slice = 0; slice < nslices; ++slice)
        {
            double t_low = timeGlobalMin + slice * dt;
            double t_high = t_low + dt;
            time->setRange(Form("slice%d", slice), t_low, t_high);
        }

        // Create a new canvas divided into 5 pads (or more, if you like)
        TCanvas *cSlices = new TCanvas();
        cSlices->Divide(nslices, 2); // creates 6 pads; we will use only 5

        // Assume nslices, dt, and timeGlobalMin are defined and cSlices is a divided canvas.
        // Assume nslices, dt, timeGlobalMin, and cSlices (a divided canvas) are defined.
        for (int slice = 0; slice < nslices; ++slice)
        {
            // Define the time slice boundaries.
            double t_low = timeGlobalMin + slice * dt;
            double t_high = t_low + dt;

            cout << "\n"
                 << endl;

            string rangeName = "slice" + to_string(slice);
            time->setRange(t_low, t_high);

            // Create a frame for the minv projection.
            RooPlot *plotMinvSlice = minv->frame(
                Title(Form("m^{2}_{inv} projection for t_{e#gamma} #in [%.02f, %.02f] ns", t_low, t_high)));

            data.plotOn(plotMinvSlice, CutRange(rangeName.c_str()));
            const double nData = data.sumEntries("", rangeName.c_str());

            // Now plot the full model.
            model.plotOn(plotMinvSlice,
                         Range("leftSide"),
                         NormRange("leftSide,rightSide"),
                         ProjectionRange(rangeName.c_str()),
                         LineColor(kMagenta + 2));

            model.plotOn(plotMinvSlice,
                         Range("rightSide"),
                         NormRange("leftSide,rightSide"),
                         ProjectionRange(rangeName.c_str()),
                         LineColor(kMagenta + 2));

            model.plotOn(plotMinvSlice,
                         Range("fullRange"),
                         NormRange("leftSide,rightSide"),
                         Components((model_name + "_rmd").c_str()),
                         ProjectionRange(rangeName.c_str()),
                         LineColor(kRed));

            model.plotOn(plotMinvSlice,
                         Range("fullRange"),
                         NormRange("leftSide,rightSide"),
                         Components((model_name + "_acc").c_str()),
                         ProjectionRange(rangeName.c_str()),
                         LineColor(kBlue));

            model.plotOn(plotMinvSlice,
                         Range("fullRange"),
                         NormRange("leftSide,rightSide"),
                         ProjectionRange(rangeName.c_str()),
                         LineColor(kMagenta + 2),
                         LineStyle(kDashed));

            // Also overlay the data (which is now restricted to the current slice).
            data.plotOn(plotMinvSlice, CutRange(rangeName.c_str()));

            // Create pull histogram
            RooPlot *pull = minv->frame(Title(Form("Pull distribution for slice %d", slice)));
            RooHist *hpull = plotMinvSlice->pullHist();
            RooHist *newHist = RemovePointsInRange(hpull, -84, 121);
            pull->SetTitle("");
            pull->addPlotable(newHist, "p");
            pull->GetYaxis()->SetTitle("pull histogram");
            pull->GetYaxis()->SetRangeUser(-3.5, +3.5);

            cSlices->cd(slice + 1 + nslices);
            pull->Draw();
            TLine *line = new TLine(minv->getMin(), 0, minv->getMax(), 0);
            line->SetLineColor(kMagenta + 2);
            line->SetLineStyle(kDashed);
            line->Draw("same");

            // Calculate chi^2/ndof
            vector<double> chi2_ndof = calculateChi2NDOF(newHist, model.getParameters(RooArgSet(*minv, *time))->getSize());
            // Add a TBox to display chi^2/ndof
            TPaveText *pt = new TPaveText(0.40, 0.1, 0.60, 0.2, "NDC");
            pt->AddText(Form("#chi^{2}/#nu = %.02f/%0.f = %.02f", chi2_ndof[0], chi2_ndof[1], chi2_ndof[0] / chi2_ndof[1]));
            pt->SetFillColor(0);
            pt->SetTextColor(kBlack);
            pt->SetTextSize(0.05);
            pt->Draw("same");

            // Draw the plot.
            cSlices->cd(slice + 1);
            plotMinvSlice->Draw();
        }
        cSlices->Update();

        // same structure for minv slices 
        /*
        // Save the global mass range (assumes minv->getMin() and minv->getMax() return the full range)
        double minvGlobalMin = minv->getMin();
        double minvGlobalMax = minv->getMax();
        int nslices_mass = 3;
        double dm = (minvGlobalMax - minvGlobalMin) / double(nslices_mass);

        time->setRange(timeGlobalMin, timeGlobalMax);

        // Create ranges for the mass slices avoiding the range -84 to 121
        for (int slice = 0; slice < nslices_mass; ++slice)
        {
            double m_low = minvGlobalMin + slice * dm;
            double m_high = m_low + dm;

            // Adjust the ranges to avoid -84 to 121
            if (m_low < -84 && m_high > -84 && m_high <= 121)
            {
                m_high = -84;
            }
            else if (m_low >= -84 && m_low < 121 && m_high > 121)
            {
                m_low = 121;
            }
            else if (m_low < -84 && m_high > 121)
            {
                m_high = m_low + dm + (121 - (-84));
            }

            minv->setRange(Form("mass_slice%d", slice), m_low, m_high);
        }

        // Create a new canvas divided into 5 pads (or more, if you like)
        TCanvas *cMassSlices = new TCanvas();
        cMassSlices->Divide(nslices_mass, 2); // creates 6 pads; we will use only 5

        // Assume nslices_mass, dm, and minvGlobalMin are defined and cMassSlices is a divided canvas.
        for (int slice = 0; slice < nslices_mass; ++slice)
        {
            // Define the mass slice boundaries.
            double m_low = minvGlobalMin + slice * dm;
            double m_high = m_low + dm;

            cout << "\n"
                 << endl;

            string rangeName = "mass_slice" + to_string(slice);
            minv->setRange(m_low, m_high);

            // Create a frame for the time projection.
            RooPlot *plotTimeSlice = time->frame(
                Title(Form("t_{e#gamma} projection for m^{2}_{inv} #in [%.02f, %.02f]", m_low, m_high)));

            data.plotOn(plotTimeSlice, CutRange(rangeName.c_str()));

            // Now plot the full model.
            model.plotOn(plotTimeSlice,
                         Components(pdf_gaus->GetName()),
                         ProjectionRange(rangeName.c_str()),
                         LineColor(kRed));

            model.plotOn(plotTimeSlice,
                         Components(pdf_cte->GetName()),
                         ProjectionRange(rangeName.c_str()),
                         LineColor(kBlue));

            model.plotOn(plotTimeSlice,
                         ProjectionRange(rangeName.c_str()),
                         LineColor(kMagenta + 2),
                         LineStyle(kDashed));

            // Also overlay the data (which is now restricted to the current slice).
            data.plotOn(plotTimeSlice, CutRange(rangeName.c_str()));

            // Create pull histogram
            RooPlot *pullMass = time->frame(Title(Form("Pull distribution for mass slice %d", slice)));
            RooHist *hpullMass = plotTimeSlice->pullHist();
            RooHist *newHistMass = RemovePointsInRange(hpullMass, 0, 0);
            pullMass->SetTitle("");
            pullMass->addPlotable(newHistMass, "p");
            pullMass->GetYaxis()->SetTitle("pull histogram");
            pullMass->GetYaxis()->SetRangeUser(-3.5, +3.5);

            cMassSlices->cd(slice + 1 + nslices_mass);
            pullMass->Draw();
            TLine *lineMass = new TLine(time->getMin(), 0, time->getMax(), 0);
            lineMass->SetLineColor(kMagenta + 2);
            lineMass->SetLineStyle(kDashed);
            lineMass->Draw("same");

            // Calculate chi^2/ndof
            vector<double> chi2_ndof_mass = calculateChi2NDOF(newHistMass, model.getParameters(RooArgSet(*minv, *time))->getSize());
            // Add a TBox to display chi^2/ndof
            TPaveText *ptMass = new TPaveText(0.40, 0.1, 0.60, 0.2, "NDC");
            ptMass->AddText(Form("#chi^{2}/#nu = %.02f/%0.f = %.02f", chi2_ndof_mass[0], chi2_ndof_mass[1], chi2_ndof_mass[0] / chi2_ndof_mass[1]));
            ptMass->SetFillColor(0);
            ptMass->SetTextColor(kBlack);
            ptMass->SetTextSize(0.05);
            ptMass->Draw("same");

            // Draw the plot.
            cMassSlices->cd(slice + 1);
            plotTimeSlice->Draw();
        }
        cMassSlices->Update();
        */
    }
}