#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 <string>
#include <iostream>

const int nBins = 200;

template <typename TVar>
class DrawVar
{
private:
    std::string _identifier;
    TVar _min;
    TVar _max;
    TVar _value = TVar{};
    TH1 *_histogram = nullptr;

public:
    DrawVar(std::string identifier, TVar min, TVar max)
        : _identifier{identifier}, _min{min}, _max{max}
    {
    }

    std::string getId()
    {
        return _identifier;
    }

    const char *getIdStr()
    {
        return _identifier.c_str();
    }

    TVar getMin()
    {
        return _min;
    }

    TVar getMax()
    {
        return _max;
    }

    TVar *getValuePointer()
    {
        return &_value;
    }

    TVar getValue()
    {
        return _value;
    }

    TH1 *getHist()
    {
        return _histogram;
    }

    void setHist(TH1D *hist)
    {
        _histogram = hist;
    }
};

int general_cuts()
{
    // just some plotting options
    gROOT->SetStyle("Plain");
    TPad foo;
    const int NRGBs = 5;
    const int NCont = 250;
    double stops[NRGBs] = {0.00, 0.34, 0.61, 0.84, 1.00};
    double red[NRGBs] = {0.00, 0.00, 0.87, 1.00, 0.51};
    double green[NRGBs] = {0.00, 0.81, 1.00, 0.20, 0.00};
    double blue[NRGBs] = {0.51, 1.00, 0.12, 0.00, 0.00};
    TColor::CreateGradientColorTable(NRGBs, stops, red, green, blue, NCont);
    gStyle->SetNumberContours(NCont);
    gStyle->SetLabelFont(132, "xyz");
    gStyle->SetTitleFont(132, "xyz");
    gStyle->SetLegendFont(132);
    gStyle->SetStatFont(132);
    gStyle->SetEndErrorSize(10.0);
    gStyle->SetOptStat(0);
    gStyle->SetOptFit(0);

    std::vector<DrawVar<Float_t>> float_vars{
        DrawVar<Float_t>("B0_PT", 0., 20000.)};
    std::vector<DrawVar<Double_t>> double_vars{

    };

    // files to load
    std::vector<std::string> filenames =
        {
            "./B0ToKpPimMuMu_Collision23_Beam6800GeV-VeloClosed-MagDown-Excl-UT_RealData_Sprucing23r1_90000000_RD.root"};

    TChain *chain = new TChain("SpruceRD_B0ToKpPimMuMu/DecayTree");
    for (unsigned int i = 0; i < filenames.size(); i++)
    {
        chain->Add(filenames.at(i).c_str());
    }

    Double_t B0_M;
    chain->SetBranchAddress("B0_M", &B0_M);
    TH1D *h1_B0_M = new TH1D("h1_B0_M", "B+ Mass", nBins, 4500, 7500);

    for (size_t i = 0; i < float_vars.size(); i++)
    {
        chain->SetBranchAddress(float_vars[i].getIdStr(), float_vars[i].getValuePointer());
        float_vars[i].setHist(new TH1D(TString::Format("h1_%s", float_vars[i].getIdStr()).Data(), float_vars[i].getIdStr(), nBins, float_vars[i].getMin(), float_vars[i].getMax()));
    }

    for (size_t i = 0; i < double_vars.size(); i++)
    {
        chain->SetBranchAddress(double_vars[i].getIdStr(), double_vars[i].getValuePointer());
        double_vars[i].setHist(new TH1D(TString::Format("h1_%s", double_vars[i].getIdStr()).Data(), double_vars[i].getIdStr(), nBins, double_vars[i].getMin(), double_vars[i].getMax()));
    }

    unsigned int entries = chain->GetEntries();
    // loop over all entries in the tree
    for (unsigned int i = 0; i < entries; i++)
    {
        chain->GetEntry(i);
        h1_B0_M->Fill(B0_M);

        for (size_t i = 0; i < float_vars.size(); i++)
        {
            float_vars[i].getHist()->Fill(float_vars[i].getValue());
        }

        for (size_t i = 0; i < double_vars.size(); i++)
        {
            double_vars[i].getHist()->Fill(double_vars[i].getValue());
        }
    }

    TCanvas *c1 = new TCanvas("c1", "Canvas 1", 0, 0, 800, 600);

    h1_B0_M->Draw();

    c1->Draw();

    if (float_vars.size() > 0)
    {
        TCanvas *c2 = new TCanvas("c2", "Canvas 2", 0, 0, 1200, 600);
        c2->Divide(4, 2);

        for (size_t i = 0; i < float_vars.size(); i++)
        {
            c2->cd(i + 1);
            float_vars[i].getHist()->Draw();
        }

        c2->Draw();
    }

    if (double_vars.size() > 0)
    {
        TCanvas *c3 = new TCanvas("c3", "Canvas 3", 0, 0, 1200, 600);
        c3->Divide(4, 2);

        for (size_t i = 0; i < double_vars.size(); i++)
        {
            c3->cd(i + 1);
            double_vars[i].getHist()->Draw();
        }

        c3->Draw();
    }

    return 0;
}