#include "TH1D.h"
#include "TH2D.h"
#include "THStack.h"
#include "TGraph.h"
#include "TTree.h"
#include "TChain.h"
#include "TFile.h"
#include "TCanvas.h"
#include "TROOT.h"
#include "TStyle.h"
#include "TColor.h"
#include "TLorentzVector.h"
#include "TRandom3.h"
#include "TLorentzVector.h"

#include "RooDataHist.h"
#include "RooRealVar.h"
#include "RooPlot.h"
#include "RooGaussian.h"
#include "RooExponential.h"
#include "RooRealConstant.h"
#include "RooAddPdf.h"
#include "RooFitResult.h"
#include "RooProduct.h"
#include "RooCrystalBall.h"
#include "RooBreitWigner.h"

#include <string>
#include <iostream>
#include <cmath>

const int nBins = 70;
const Double_t MASS_HIST_MIN = 4000.;
const Double_t MASS_HIST_MAX = 8500.;
const Double_t MASS_HIST_FIT_MIN = 5100.;
const Double_t MASS_HIST_FIT_MAX = 6000.;
const Double_t J_PSI_MASS = 3096.916;
const std::string SAVE_PATH = "/work/pfeiffer/inclusive_detached_dilepton/status_report";

struct AnalysisOutput {
    std::string title;
    std::string name;
    std::string root_file;
    std::string root_file_tree;
    std::string B_name;
    std::string paper_yield;

    const std::string suf(std::string input) {
        return TString::Format("%s_%s", input.c_str(), name.c_str()).Data();
    }

    const std::string title_suf(std::string input) {
        return TString::Format("%s [%s]", input.c_str(), title.c_str()).Data();
    }
};

std::vector<RooPlot*> CreateRooFit(TH1D *hist, AnalysisOutput ana);

bool inRange(double value, double center, double low_intvl, double up_intvl) {
    return center - low_intvl < value && value < center + up_intvl;
}

bool inRange(double value, double center, double intvl) {
    return inRange(value, center, intvl, intvl);
}

void savePDF(TObject* o, const char* decay, const char* name, Option_t* opt = "") {
    auto cname = TString::Format("%s_%s", decay, name);
    std::cout << " ----- " << cname.Data() << " - " << decay << " - " << name << std::endl;
    auto c = new TCanvas(cname.Data(), cname.Data(), 0, 0, 800, 800);
    o->Draw(opt);
    c->SaveAs(TString::Format("%s/%s.jpg", SAVE_PATH.c_str(), cname.Data()));
    // c->Draw();
}

void savePDF(TObject* o1, TObject* o2, const char* decay, const char* name) {
    auto cname = TString::Format("%s_%s", decay, name);
    std::cout << " ----- " << cname.Data() << " - " << decay << " - " << name << std::endl;
    auto c = new TCanvas(cname.Data(), cname.Data(), 0, 0, 800, 800);
    auto p2 = new TPad(TString::Format("%s_p2", name), "Lower Pad", 0., 0., 1., 0.3);
    p2->Draw();
    p2->SetTopMargin(0.001);
    p2->SetBottomMargin(0.3);
    p2->SetGrid();

    auto *p1 = new TPad(TString::Format("%s_p1", name), "Upper Pad", 0., 0.32, 1., 1.);
    p1->Draw();
    p1->SetBottomMargin(0.001);
    p1->cd();

    o1->Draw();

    p2->cd();

    o2->Draw();
    c->SaveAs(TString::Format("%s/%s.jpg", SAVE_PATH.c_str(), cname.Data()));
}

RooHelpers::LocalChangeMsgLevel changeMsgLvl(RooFit::WARNING);

int status_report_plots()
{
    // gROOT->ProcessLine(".x /work/pfeiffer/lhcbStyle.C");

    std::vector<AnalysisOutput> anas {
        // AnalysisOutput{"SpruceRD_BuToHpMuMu", "BuToKpMuMu_incl", "/auto/data/pfeiffer/inclusive_detached_dilepton/data_samples/BuToKpMuMu_Collision23_Beam6800GeV-VeloClosed-MagDown-Excl-UT_RealData_Sprucing23r1_90000000_RD.root", "SpruceRD_BuToKpMuMu/DecayTree", "Bplus", ""},
        // AnalysisOutput{"BuToKpMuMu Inclusive, AnaProd NTuple", "BuToKpMuMu_incl_ap", "/auto/data/pfeiffer/inclusive_detached_dilepton/data_samples/spruce_magdown_2023_v0r1_tuple_90000000_2023_v0r0p6288631.root", "BuToHpMuMu/DecayTree", "B", "1395.27 +/- 58"},
        // AnalysisOutput{"Hlt2RD_BuToKpMuMu_2023", "BuToKpMuMu_excl", "/auto/data/pfeiffer/inclusive_detached_dilepton/data_samples/BuToKpMuMu_Collision23_Beam6800GeV-VeloClosed-MagDown-Excl-UT_RealData_SprucingPass23r1_94000000_RD.root", "Hlt2RD_BuToKpMuMu_2023/DecayTree", "Bplus", "803.769 +/- 34" },
        // AnalysisOutput{"Hlt2RD_", "BuToKpMuMu_excl_ap23", "/auto/data/pfeiffer/inclusive_detached_dilepton/data_samples/turbopass_magdown_2023_v1_tuple_94000000_2023_v0r0p6201764.root", "BuToKpMuMu23/DecayTree", "B", "803.769 +/- 34"},
        // AnalysisOutput{"SpruceRD_B0ToHpHmMuMu", "B0ToKpPimMuMu_incl", "/auto/data/pfeiffer/inclusive_detached_dilepton/data_samples/B0ToKpPimMuMu_Collision23_Beam6800GeV-VeloClosed-MagDown-Excl-UT_RealData_Sprucing23r1_90000000_RD.root", "SpruceRD_B0ToKpPimMuMu/DecayTree", "B0", "" },
        // AnalysisOutput{"B0ToKpPimMuMu_ap", "/auto/data/pfeiffer/inclusive_detached_dilepton/data_samples/spruce_magdown_2023_v0r1_tuple_90000000_2023_v0r0p6288631.root", "B0ToHpHmMuMu/DecayTree", "B0", ""},
        AnalysisOutput{"Hlt2RD_B0ToKpPimMuMu_2023", "B0ToKpPimMuMu_excl", "/auto/data/pfeiffer/inclusive_detached_dilepton/data_samples/B0ToKpPimMuMu_Collision23_Beam6800GeV-VeloClosed-MagDown-Excl-UT_RealData_SprucingPass23r1_94000000_RD.root", "Hlt2RD_B0ToKpPimMuMu_2023/DecayTree", "B0", "" },
        AnalysisOutput{"Hlt2RD_B0ToKpPimMuMu_2023 AP", "B0ToKpPimMuMu_excl_ap23", "/auto/data/pfeiffer/inclusive_detached_dilepton/data_samples/turbopass_magdown_2023_v1_tuple_94000000_2023_v0r0p6201764.root", "B0ToKpPimMuMu23/DecayTree", "B0", ""},
        // AnalysisOutput{"BuToJpsiKplus_det", "/auto/data/pfeiffer/inclusive_detached_dilepton/data_samples/BuToJpsiKplus_Detached_Collision23_Beam6800GeV-VeloClosed-MagDown-Excl-UT_RealData_SprucingPass23r1_94000000_B2CC.root", "Hlt2B2CC_BuToJpsiKplus_JpsiToMuMu_Detached/DecayTree", "Bplus", "unknown" },
    };

    for (size_t a = 0; a < anas.size(); a++)
    {
        auto ana = anas[a];

        TChain *data_chain = new TChain(ana.root_file_tree.c_str());
        data_chain->Add(ana.root_file.c_str());

        Double_t B_M,   // 4400 -> 8200
                Jpsi_M; // 200 -> 6600

        Float_t muplus_PX,    // 0 -> 30 * 10^3
                muplus_PY, // 0 -> 7000
                muplus_PZ,  // 0 -> 3
                muplus_ENERGY,         // 0 -> 26 * 10^3
                muminus_PX, // 0 -> 7
                muminus_PY,        // 0 -> 11 * 10^3
                muminus_PZ, // 0 -> 7
                muminus_ENERGY;        // 0 -> 11 * 10^3

        data_chain->SetBranchAddress(TString::Format("%s_M", ana.B_name.c_str()).Data(), &B_M);
        data_chain->SetBranchAddress("Jpsi_M", &Jpsi_M);

        // data_chain->SetBranchAddress("muplus_PX", &muplus_PX);
        // data_chain->SetBranchAddress("muplus_PY", &muplus_PY);
        // data_chain->SetBranchAddress("muplus_PZ", &muplus_PZ);
        // data_chain->SetBranchAddress("muplus_ENERGY", &muplus_ENERGY);
        // data_chain->SetBranchAddress("muminus_PX", &muminus_PX);
        // data_chain->SetBranchAddress("muminus_PY", &muminus_PY);
        // data_chain->SetBranchAddress("muminus_PZ", &muminus_PZ);
        // data_chain->SetBranchAddress("muminus_ENERGY", &muminus_ENERGY);

        TH1D *h1_B_M = new TH1D(ana.suf("h1_B_M").c_str(), ana.title_suf("B Mass").c_str(), nBins, MASS_HIST_MIN, MASS_HIST_MAX);
        TH1D *h1_B_M_JPsi_cut = new TH1D(ana.suf("h1_B_M_JPsi_cut").c_str(), ana.title_suf("B Mass").c_str(), nBins, MASS_HIST_FIT_MIN, MASS_HIST_FIT_MAX);
        TH1D *h1_Jpsi_M = new TH1D(ana.suf("h1_Jpsi_M").c_str(), ana.title_suf("J/#psi Mass").c_str(), nBins, 200., 5500.);
        TH2D *h2_B_M_Jpsi_M = new TH2D(ana.suf("h2_B_M_Jpsi_M").c_str(), ana.title_suf("B Mass vs. J/#psi Mass").c_str(), nBins, MASS_HIST_MIN, MASS_HIST_MAX, nBins, 200., 5500.);


        unsigned int entries = data_chain->GetEntries();

        for (unsigned int i = 0; i < entries; i++)
        {
            if (i % 10000 == 0)
            {
                std::cout << "[" << ana.name << "] Processing event: " << i << " (" << TString::Format("%.2f", ((double)i / (double)entries) * 100.) << "%)" << std::endl;
            }

            // TLorentzVector v4_muplus(muplus_PX, muplus_PY, muplus_PZ, muplus_ENERGY);
            // TLorentzVector v4_muminus(muminus_PX, muminus_PY, muminus_PZ, muminus_ENERGY);
            // Double_t mumu_inv_mass = (v4_muplus + v4_muminus).M();

            data_chain->GetEntry(i);
            h1_B_M->Fill(B_M);
            h1_Jpsi_M->Fill(Jpsi_M);
            h2_B_M_Jpsi_M->Fill(B_M, Jpsi_M);
            // std::cout << mumu_inv_mass << "-" << Jpsi_M << "=" << TMath::Abs(mumu_inv_mass-Jpsi_M) << std::endl;

            if (TMath::Abs(Jpsi_M - J_PSI_MASS) < 100.)
            {
                h1_B_M_JPsi_cut->Fill(B_M);
            }
        }

        h1_B_M->SetMinimum(0);
        h1_B_M_JPsi_cut->SetMinimum(0); 
        h1_Jpsi_M->SetMinimum(0);
        h2_B_M_Jpsi_M->SetMinimum(0);

        auto fitRes = CreateRooFit(h1_B_M_JPsi_cut, ana);

        savePDF(h1_B_M, ana.name.c_str(), "B_M_uncut");
        savePDF(h1_Jpsi_M, ana.name.c_str(), "JPsi_M_uncut");
        savePDF(h1_B_M_JPsi_cut, ana.name.c_str(), "B_M_JPsi_cut");
        savePDF(fitRes[0], fitRes[1], ana.name.c_str(), "B_mass_JPsi_cut_fit");
    }

    return 0;
}


std::vector<RooPlot*> CreateRooFit(TH1D *hist, AnalysisOutput ana)
{
    RooRealVar roo_var_mass(ana.suf("roo_var_mass").c_str(), "B Mass Variable", MASS_HIST_FIT_MIN, MASS_HIST_FIT_MAX);
    roo_var_mass.setRange("fitting_range", MASS_HIST_FIT_MIN, MASS_HIST_FIT_MAX);

    std::string hist_name = ana.suf("roohist_B_M");
    RooDataHist roohist_B_M(hist_name.c_str(), "B Mass Histogram", roo_var_mass, RooFit::Import(*hist));

    RooRealVar roo_sig_bw_mean(ana.suf("roo_sig_bw_mean").c_str(), "Mass BW Mean", 5250., 5100., 5400.);
    RooRealVar roo_sig_bw_with(ana.suf("roo_sig_bw_with").c_str(), "Mass BW Width", 20., 0., 50.);

    RooBreitWigner roo_sig_bw(ana.suf("roo_sig_bw").c_str(), "B Signal Breit Wigner", roo_var_mass, roo_sig_bw_mean, roo_sig_bw_with);

    RooRealVar roo_bkg_exp_c(ana.suf("roo_bkg_exp_c").c_str(), "Background C", -0.001145, -0.00199, -0.00100);
    RooExponential roo_bkg_exp(ana.suf("roo_bkg_exp").c_str(), "B Mass Background Exp", roo_var_mass, roo_bkg_exp_c);

    RooRealVar roo_var_mass_sig_yield(ana.suf("roo_var_mass_sig_yield").c_str(), "B Mass Sig Yield", 0., hist->GetEntries());
    RooRealVar roo_var_mass_bkg_yield(ana.suf("roo_var_mass_bkg_yield").c_str(), "B Mass Bckg Yield", 0., hist->GetEntries());

    std::string pdf_name = ana.suf("roo_pdf_sig_plus_bkg");
    RooAddPdf roo_pdf_sig_plus_bkg(pdf_name.c_str(), "Sig + Bkg PDF",
                                   RooArgList(roo_sig_bw, roo_bkg_exp),
                                   RooArgList(roo_var_mass_sig_yield, roo_var_mass_bkg_yield));

    RooPlot *roo_frame_mass = roo_var_mass.frame(RooFit::Title(TString::Format("%s [%s]", ana.name.c_str(), ana.title.c_str()).Data()));
    roohist_B_M.plotOn(roo_frame_mass, RooFit::Binning(nBins), RooFit::Name(hist_name.c_str()));

    RooFitResult *fitres = roo_pdf_sig_plus_bkg.fitTo(roohist_B_M, RooFit::Save(), RooFit::PrintLevel(1), RooFit::Range("fitting_range"));

    roo_pdf_sig_plus_bkg.plotOn(roo_frame_mass, RooFit::LineColor(kRed), RooFit::LineStyle(kSolid), RooFit::Range("fitting_range"), RooFit::Name(pdf_name.c_str()));
    roo_pdf_sig_plus_bkg.plotOn(roo_frame_mass, RooFit::Components(RooArgSet(roo_bkg_exp)), RooFit::LineColor(kBlue), RooFit::LineStyle(kDashed), RooFit::Range("fitting_range"));
    roo_pdf_sig_plus_bkg.plotOn(roo_frame_mass, RooFit::Components(RooArgSet(roo_sig_bw)), RooFit::FillStyle(3244), RooFit::LineColor(kRed-7), RooFit::LineStyle(kDashed), RooFit::Range("fitting_range"));

    roo_pdf_sig_plus_bkg.paramOn(roo_frame_mass, RooFit::Layout(0.40, 0.99, 0.90));
    roo_frame_mass->getAttText()->SetTextSize(0.024);

    RooPlot *roo_frame_pull = roo_var_mass.frame(RooFit::Title("Pull Distribution"));
    roo_frame_pull->addPlotable(roo_frame_mass->pullHist(hist_name.c_str(), pdf_name.c_str()), "P");

    return std::vector<RooPlot*> { roo_frame_mass, roo_frame_pull };
}