// VBS
#include "vbsvvhjets/cuts.h"
// RAPIDO
#include "hepcli.h"
#include "looper.h"
// ROOT
#include "TString.h"
#include "TObject.h"
#include "Math/VectorUtil.h"
#include "TH1.h"
// NanoCORE
#include "Nano.h"
#include "Config.h"
#include "ElectronSelections.h"
#include "MuonSelections.h"
#include "tqdm.h"

struct TriggerEff
{
    TString name;
    TH1F* numer;
    TH1F* denom;

    TriggerEff(TString trigger_name, int n_bins, int low, int high)
    {
        name = trigger_name;
        numer = new TH1F(name+"_eff_numer", "", n_bins, low, high);
        denom = new TH1F(name+"_eff_denom", "", n_bins, low, high);
        numer->Sumw2();
        denom->Sumw2();
    };

    void fillNumer(float x, float weight = 1.f) { numer->Fill(x, weight); }

    void fillDenom(float x, float weight = 1.f) { denom->Fill(x, weight); }

    TH1F* getEff()
    {
        TH1F* eff = (TH1F*)numer->Clone(name+"_eff");
        eff->Divide(eff, denom, 1, 1, "B");
        return eff;
    };

    void writeHists()
    {
        numer->Write();
        denom->Write();

        TH1F* eff = getEff();
        eff->Write();
    };
};

bool passPrescaled()
{
    bool passed = false;
    switch (nt.year())
    {
    case (2016):
        try { passed = (passed || nt.HLT_PFHT650()); }
        catch (const runtime_error& error) { /* do nothing */ }
        break;
    case (2017):
        try { passed = (passed || nt.HLT_PFHT890()); }
        catch (const runtime_error& error) { /* do nothing */ }
        break;
    case (2018):
        try { passed = (passed || nt.HLT_PFHT890()); }
        catch (const runtime_error& error) { /* do nothing */ }
        break;
    }
    return passed;
}

int main(int argc, char** argv) 
{
    // CLI
    HEPCLI cli = HEPCLI(argc, argv);

    // Initialize Looper
    Looper looper = Looper(cli);

    // Initialize TFile
    TFile* tfile = new TFile(TString(cli.output_dir+"/"+cli.output_name+".root"), "RECREATE");

    // Initialize histograms
    TriggerEff hlt_ak4_ht_eff = TriggerEff("HLT_AK4_HT", 60, 0, 3000);

    // Run looper
    tqdm bar;
    looper.run(
        [&](TTree* ttree) 
        { 
            nt.Init(ttree); 
        },
        [&](int entry) 
        {
            if (cli.debug && looper.n_events_processed == 10000) { looper.stop(); }
            else
            {
                nt.GetEntry(entry);

                if (!passPrescaled()) { return; }

                float ak4_ht = 0;
                for (unsigned int ak4_i = 0; ak4_i < nt.nJet(); ++ak4_i)
                {
                    ak4_ht += nt.Jet_pt().at(ak4_i);
                }

                bool passed = false;
                switch (nt.year())
                {
                case (2016):
                    try { passed = (passed || nt.HLT_PFHT800()); }
                    catch (const runtime_error& error) { /* do nothing */ }
                    try { passed = (passed || nt.HLT_PFHT900()); }
                    catch (const runtime_error& error) { /* do nothing */ }
                    break;
                case (2017):
                    try { passed = (passed || nt.HLT_PFHT1050()); }
                    catch (const runtime_error& error) { /* do nothing */ }
                    break;
                case (2018):
                    try { passed = (passed || nt.HLT_PFHT1050()); }
                    catch (const runtime_error& error) { /* do nothing */ }
                    break;
                }

                if (passed) { hlt_ak4_ht_eff.fillNumer(ak4_ht); }
                hlt_ak4_ht_eff.fillDenom(ak4_ht);

                bar.progress(looper.n_events_processed, looper.n_events_total);
            }
        }
    );

    // Wrap up
    tfile->cd();
    hlt_ak4_ht_eff.writeHists();
    tfile->Close();

    return 0;
}