#include "RooRealVar.h"
#include "RooDataHist.h"
#include "RooGaussian.h"
#include "RooExponential.h"
#include "RooProdPdf.h"
#include "RooAddPdf.h"
#include "RooFitResult.h"
#include "RooPlot.h"
#include "TH2.h"
#include "TCanvas.h"
#include "TFile.h"
#include "TMath.h"

using namespace RooFit;

void Fit2DConstant_Example2() {
    // Read in data histograms from file
    TFile* file = TFile::Open("histograms.root");
    TH3D* data1_3D = static_cast<TH3D*>(file->Get("data1")); 
    data1_3D->SetTitle("Data 1 Histogram");
    TH3D* data2_3D = static_cast<TH3D*>(file->Get("data2")); 
    data2_3D->SetTitle("Data 2 Histogram");
    TH3D* signal_3D = static_cast<TH3D*>(file->Get("signal")); 
    signal_3D->SetTitle("Signal Histogram");
    TH3D* bg1_constant_3D = static_cast<TH3D*>(file->Get("background1_constant")); 
    bg1_constant_3D->SetTitle("Background 1 Constant Histogram");
    TH1D* bg1_scaling_1D = static_cast<TH1D*>(file->Get("background1_scaling")); 
    bg1_scaling_1D->SetTitle("Background 1 Scaling Histogram");
    TH3D* bg2_3D = static_cast<TH3D*>(file->Get("background2")); 
    bg2_3D->SetTitle("Background 2 Histogram");
    TH3D* bg3_3D = static_cast<TH3D*>(file->Get("background3"));
    bg3_3D->SetTitle("Background 3 Histogram");
    TH3D* bg4_3D = static_cast<TH3D*>(file->Get("background4")); 
    bg4_3D->SetTitle("Background 4 Histogram");

    // Create combined data histogram
    TH3D* data_3D = static_cast<TH3D*>(data1_3D->Clone("data_combined"));
    data_3D->Sumw2();
    data_3D->Add(data2_3D);

    // Scale background 1 histogram 
    bg1_constant_3D->Sumw2();
    double areaDensityData = 5.053e-8;
    double areaDensityBackground1 = 4.237e-8; 
    double areaDensityFactor = areaDensityData / areaDensityBackground1;
    bg1_constant_3D->Scale(areaDensityFactor);
    for (int i = 1; i <= bg1_constant_3D->GetNbinsX(); ++i) {
        double scalingFactor = bg1_scaling_1D->GetBinContent(i);
        for (int j = 1; j <= bg1_constant_3D->GetNbinsY(); ++j) {
            for (int k = 1; k <= bg1_constant_3D->GetNbinsZ(); ++k) {
                double content = bg1_constant_3D->GetBinContent(i, j, k);
                bg1_constant_3D->SetBinContent(i, j, k, content * scalingFactor);
            }
        }
    }

    // Helper lambda for 3D → 2D projection
    auto project3DTo2D = [](TH3* hist, const char* name, int xBinLow, int xBinHigh) {
        hist->GetXaxis()->SetRange(xBinLow, xBinHigh);
        TH2D* hist2D = static_cast<TH2D*>(hist->Project3D("zy"));
        hist->GetXaxis()->SetRange(1, hist->GetNbinsX()); // reset range
        hist2D->SetName(name);
        return hist2D;
    };

    int xBinLow = 20;
    int xBinHigh = 30;

    TH2D* data_2D = project3DTo2D(data_3D, "data_2D", xBinLow, xBinHigh); 
    data_2D->SetTitle("Data 2D Histogram");
    TH2D* signal_2D = project3DTo2D(signal_3D, "signal_2D", xBinLow, xBinHigh); 
    signal_2D->SetTitle("Signal 2D Histogram");
    TH2D* bg1_constant_2D = project3DTo2D(bg1_constant_3D, "bg1_constant_2D", xBinLow, xBinHigh); 
    bg1_constant_2D->SetTitle("Background 1 Constant 2D Histogram");
    TH2D* bg2_2D = project3DTo2D(bg2_3D, "bg2_2D", xBinLow, xBinHigh); 
    bg2_2D->SetTitle("Background 2 2D Histogram");
    TH2D* bg3_2D = project3DTo2D(bg3_3D, "bg3_2D", xBinLow, xBinHigh); 
    bg3_2D->SetTitle("Background 3 2D Histogram");
    TH2D* bg4_2D = project3DTo2D(bg4_3D, "bg4_2D", xBinLow, xBinHigh); 
    bg4_2D->SetTitle("Background 4 2D Histogram");

    // Add small number so no empty bins in RooHistPdf (otherwise fit fails)
    double smallNumber = 1e-6;
    auto addSmallNumber = [smallNumber](TH2* hist) {
        for (int i = 1; i <= hist->GetNbinsX(); ++i) {
            for (int j = 1; j <= hist->GetNbinsY(); ++j) {
                if (hist->GetBinContent(i, j) < smallNumber) {
                    hist->SetBinContent(i, j, smallNumber);
                }
            }
        }
    };
    addSmallNumber(data_2D);
    addSmallNumber(signal_2D);
    addSmallNumber(bg1_constant_2D);
    addSmallNumber(bg2_2D);
    addSmallNumber(bg3_2D);
    addSmallNumber(bg4_2D);

    // Define observables
    RooRealVar x("x", "Observable X", data_2D->GetXaxis()->GetXmin(), data_2D->GetXaxis()->GetXmax());
    RooRealVar y("y", "Observable Y", data_2D->GetYaxis()->GetXmin(), data_2D->GetYaxis()->GetXmax());
    x.setBins(data_2D->GetXaxis()->GetNbins());
    y.setBins(data_2D->GetYaxis()->GetNbins());

    // Wrap in RooDataHists
    RooDataHist data("data_2D", "Data 2D", RooArgSet(x, y), data_2D);
    RooDataHist signal_data("signal_2D", "Signal Data 2D", RooArgSet(x, y), signal_2D);
    RooDataHist bg1_constant_data("bg1_constant_2D", "BG1 Const 2D", RooArgSet(x, y), bg1_constant_2D);
    RooDataHist bg2_data("bg2_2D", "BG2 2D", RooArgSet(x, y), bg2_2D);
    RooDataHist bg3_data("bg3_2D", "BG3 2D", RooArgSet(x, y), bg3_2D);
    RooDataHist bg4_data("bg4_2D", "BG4 2D", RooArgSet(x, y), bg4_2D);

    // Create PDFs
    RooHistPdf signal_pdf("signal_pdf", "Signal PDF", RooArgSet(x, y), signal_data);
    RooHistPdf bg1_constant_pdf("bg1_constant_pdf", "BG1 PDF", RooArgSet(x, y), bg1_constant_data);
    RooHistPdf bg2_pdf("bg2_pdf", "BG2 PDF", RooArgSet(x, y), bg2_data);
    RooHistPdf bg3_pdf("bg3_pdf", "BG3 PDF", RooArgSet(x, y), bg3_data);
    RooHistPdf bg4_pdf("bg4_pdf", "BG4 PDF", RooArgSet(x, y), bg4_data);

    int dataEntries = data_2D->Integral();
    int startYield = round(dataEntries * 0.1);

    RooRealVar nsig("nsig", "Signal yield", startYield, 0, dataEntries);
    RooConstVar nbkg1("nbkg1", "BG1 yield", bg1_constant_2D->Integral());
    RooRealVar nbkg2("nbkg2", "BG2 yield", startYield, 0, dataEntries);
    RooRealVar nbkg3("nbkg3", "BG3 yield", startYield, 0, dataEntries);
    RooRealVar nbkg4("nbkg4", "BG4 yield", startYield, 0, dataEntries);

    RooAddPdf model("model", "Total Model",
        RooArgList(signal_pdf, bg1_constant_pdf, bg2_pdf, bg3_pdf, bg4_pdf),
        RooArgList(nsig, nbkg1, nbkg2, nbkg3, nbkg4)
    );

    // Print pre-fit expectations
    RooArgSet nset{x, y};
    std::cout << "nbkg3          : " << nbkg3.getVal() << std::endl;
    std::cout << "dataset entries: " << data.numEntries() << std::endl;
    std::cout << "expected before: " << model.expectedEvents(&nset) << std::endl;

    // Fit model
    RooFitResult* fitResult = model.fitTo(data, Save(), SumW2Error(true), PrintLevel(-1), PrintEvalErrors(-1));

    std::cout << "expected after : " << model.expectedEvents(&nset) << std::endl;
    fitResult->Print("v");

    // RooCmdArg normArg = Normalization(model.expectedEvents(&nset) / data.numEntries());
    RooCmdArg normArg = Normalization(1.0); // Use 1.0 for normalization

    // Make plots
    TCanvas* c1 = new TCanvas("c1", "Canvas", 1200, 800);
    c1->Divide(4, 3);
    c1->cd(1); data_2D->Draw("COLZ");
    c1->cd(2); signal_2D->Draw("COLZ");
    c1->cd(3); bg1_constant_2D->Draw("COLZ");
    c1->cd(4); bg2_2D->Draw("COLZ");
    c1->cd(5); bg3_2D->Draw("COLZ");
    c1->cd(6); bg4_2D->Draw("COLZ");
    // Draw data in black and background in red
    c1->cd(7); 
    TH1D* data_projX = data_2D->ProjectionX("data_projX");
    data_projX->SetTitle("Projection X of Data");
    TH1D* bg1_constant_projX = bg1_constant_2D->ProjectionX("bg1_constant_projX");
    bg1_constant_projX->SetTitle("Projection X of Background 1 Constant");
    data_projX->SetLineColor(kBlack);
    bg1_constant_projX->SetLineColor(kRed);
    data_projX->Draw();
    bg1_constant_projX->Draw("SAME");
    c1->cd(8);
    TH1D* data_projY = data_2D->ProjectionY("data_projY");
    data_projY->SetTitle("Projection Y of Data");
    TH1D* bg1_constant_projY = bg1_constant_2D->ProjectionY("bg1_constant_projY");
    bg1_constant_projY->SetTitle("Projection Y of Background 1 Constant");
    data_projY->SetLineColor(kBlack);
    bg1_constant_projY->SetLineColor(kRed);
    data_projY->Draw();
    bg1_constant_projY->Draw("SAME");

    // X projection
    c1->cd(9);
    RooPlot* framex = x.frame();
    framex->SetTitle("Fit in X");
    data.plotOn(framex);
    model.plotOn(framex, normArg);
    model.plotOn(framex, Components("signal_pdf"), LineColor(kGreen), LineStyle(kDashed), normArg);
    model.plotOn(framex, Components("bg1_constant_pdf"), LineColor(kRed), LineStyle(kDashed), normArg);
    model.plotOn(framex, Components("bg2_pdf"), LineColor(kBlue), LineStyle(kDashed), normArg);
    model.plotOn(framex, Components("bg3_pdf"), LineColor(kMagenta), LineStyle(kDashed), normArg);
    model.plotOn(framex, Components("bg4_pdf"), LineColor(kCyan), LineStyle(kDashed), normArg);
    framex->Draw();

    // Y projection
    c1->cd(10);
    RooPlot* framey = y.frame();
    framey->SetTitle("Fit in Y");
    data.plotOn(framey);
    model.plotOn(framey, normArg);
    model.plotOn(framey, Components("signal_pdf"), LineColor(kGreen), LineStyle(kDashed), normArg);
    model.plotOn(framey, Components("bg1_constant_pdf"), LineColor(kRed), LineStyle(kDashed), normArg);
    model.plotOn(framey, Components("bg2_pdf"), LineColor(kBlue), LineStyle(kDashed), normArg);
    model.plotOn(framey, Components("bg3_pdf"), LineColor(kMagenta), LineStyle(kDashed), normArg);
    model.plotOn(framey, Components("bg4_pdf"), LineColor(kCyan), LineStyle(kDashed), normArg);
    framey->Draw();
}