//Helper functions
//Yes, they are in a header, because ROOT
//Renata Kopecna
#ifndef UTILS_HPP
#define UTILS_HPP

#include "GlobalFunctions.hh"
#include "Paths.hpp"
#include "MassFit.hpp"
#include <boost/algorithm/string.hpp>

//I HATE ROOT! It's just way easier for root to have functions defined in a header with cross-references.... sigh

//Forward declarations from Paths.hpp
string GetInputFile(string year, string magnet, bool preSelected, bool MC, bool ReferenceChannel, bool PHSP,  bool B0, bool K1, bool Inc, bool smallSample);
string GetBDTinputFile(string year, bool MC, bool ReferenceChannel, bool PHSP, bool KshortDecayInVelo);
string GetBDToutputFile(string year, int Run, bool MC, bool ReferenceChannel, bool PHSP, bool KshortDecayInVelo, bool UseLowQ2Range, bool reweighted);
string GetBDToutputFile(int Run, bool MC, bool ReferenceChannel, bool PHSP, bool KshortDecayInVelo, bool UseLowQ2Range, bool reweighted);
string returnFileAddress(string year, int Run, string magnet, bool Preselected, bool BDTed, bool MC, bool ReferenceChannel, bool PHSP, bool KshortDecayInVelo);

/////////////////////////////////////////////
//
//          Get basic names and trees
//
/////////////////////////////////////////////

string correct_magnet_string(string magnet){
    boost::to_lower(magnet);

    if (magnet != "both" && magnet != "up" && magnet != "down"){
        coutWarning("Wrong magnet polarity used! Setting magnet to both!");
        return magnet = "both";
    }
    else return magnet;
}

string treeName(bool MC, bool Preselected){
    if (MC && Preselected) return "DecayTreeTruthMatched";
    else if (!Preselected) return "b2KstKpi0mumuResolvedTuple/DecayTree";
    else return "DecayTree";
}

vector <string> get_magnet_vector(string magnet){
    magnet = correct_magnet_string(magnet);
    if (magnet == "both"){
        return {"down","up"};
    }
    else return {magnet};

}

TChain *get_basic_TChain(string magnet, vector<string> years, bool Preselected, bool MC, bool ReferenceChannel, bool PHSP, bool B0, bool K1, bool Inc){

    TChain* tree = new TChain(treeName(MC,Preselected).c_str());
    coutInfo("Reading data from TTree... ");

    for (auto& yr : years){
        for (auto& mag: get_magnet_vector(magnet)){
            string path = GetInputFile(yr,mag,Preselected,MC, ReferenceChannel, PHSP, B0, K1, Inc, false);
            coutDebug("Adding " + path);
            tree->Add(path.c_str()); //TODO: no acces to not-preselected inclusive sample
        }
    }

    return tree;
}

TChain *get_weighted_TChain(vector<string> years, bool MC, bool ReferenceChannel, bool PHSP, bool KshortDecayInVelo){

    TChain* tree = new TChain(treeName(MC,true).c_str());
    for (auto& yr : years){
        tree->Add(GetBDTinputFile(yr,MC,ReferenceChannel,PHSP,KshortDecayInVelo).c_str());
        coutDebug("Adding " +GetBDTinputFile(yr,MC,ReferenceChannel,PHSP,KshortDecayInVelo));
    }
    return tree;

}

TChain *get_BDT_TChain(vector<string> years, bool MC, bool ReferenceChannel, bool PHSP, bool KshortDecayInVelo, bool reweighted){
    bool UseLowQ2Range = false;

    TChain* tree = new TChain(treeName(MC,true).c_str());
    for (auto& yr : years){
        tree->Add(GetBDToutputFile(yr,getRunID(yr),MC,ReferenceChannel,PHSP,KshortDecayInVelo,UseLowQ2Range,reweighted).c_str());
        coutDebug("Adding " +GetBDToutputFile(yr,getRunID(yr),MC,ReferenceChannel,PHSP,KshortDecayInVelo,UseLowQ2Range,reweighted));
    }
    return tree;
}

TChain *get_BDT_TChain(string year, int Run, bool MC, bool ReferenceChannel, bool PHSP, bool KshortDecayInVelo, bool reweighted){
    bool UseLowQ2Range = false;

    vector<string> years;
    if (Run == 0) years.push_back(year);
    else years = yearsVector(MC,ReferenceChannel,PHSP,Run);

    TChain* tree = new TChain(treeName(MC,true).c_str());
    for (auto& yr : years){
        tree->Add(GetBDToutputFile(yr,getRunID(yr),MC,ReferenceChannel,PHSP,KshortDecayInVelo,UseLowQ2Range,reweighted).c_str());
        coutDebug("Adding " +GetBDToutputFile(yr,getRunID(yr),MC,ReferenceChannel,PHSP,KshortDecayInVelo,UseLowQ2Range,reweighted));
    }
    return tree;
}


TChain *get_FinalOutput_TChain(vector<string> years, bool MC, bool ReferenceChannel, bool PHSP, bool KshortDecayInVelo){ //TODO fix Run vs year

    TChain* tree = new TChain(treeName(MC,true).c_str());
    for (auto& yr : years){
        //tree->Add(GetFinalOutputFile(yr,1,MC,ReferenceChannel,PHSP,KshortDecayInVelo).c_str());
    }
    return tree;
}

///////////////////////////////////////////////
//
//          Cut functions
//
///////////////////////////////////////////////

string getTMcut(bool TM,bool notTM, string customTMbranch, bool gammaTM){

     string finalCut = "";
     if (TM && notTM){
         coutERROR("Cannot cut on TM == 1 && TM == 0! Returning an empty string.");
         return "";
     }
     if (TM){
         if (customTMbranch == "") finalCut.append("(" + TMbranch + "==1)");
         else                      finalCut.append("(" + customTMbranch + "==1)");
         //Check if the TM branch is not BKGCAT
         if (finalCut.find("BKGCAT") == string::npos){
             //Gamma category: 1: both passed, 2: one passed, one converted, 3: both converted, 4: one passed, one random, 5: one converted, one random, 6: both failed
             if (gammaTM) finalCut.append(" && (" + gammaTMbranch + "<4)");
             else         finalCut.append(" && (" + gammaTMbranch + "<6)");
         }
     }
     if (notTM){
         if (customTMbranch == "") finalCut.append("(" + TMbranch + "==0)");
         else                      finalCut.append("(" + customTMbranch + "==0)");
         if (finalCut.find("BKGCAT") == string::npos){
             //Gamma category: 1: both passed, 2: one passed, one converted, 3: both converted, 4: one passed, one random, 5: one converted, one random, 6: both failed
             if (gammaTM) finalCut.append(" || (" + gammaTMbranch + ">3)");
             else         finalCut.append(" || (" + gammaTMbranch + ">5)");
         }
     }
     coutTest("TM cut:" + finalCut);
     return finalCut;
}

string getJpsicut(){
   return "(8.68e6 < Q2 && Q2 < 10.09e6)"; //q2Cuts = "nDiMuonMassBin == 9";
}

string getMuMucut(){
   return "(Q2 < 0.98e6 || (1.1e6 < Q2 && Q2 < 8.0e6) || (11.0e6 < Q2 && Q2 < 12.5e6) || Q2 > 15.0e6)";
}

string getFullCut(bool UseOnlyMuMuEvents, bool UseOnlyJpsiEvents, bool SplitInQ2, bool UseLowQ2Range, bool useExtraVar, TMefficiencyClass extraVar,int nExtraBin, double TMVAcut){
    string finalCut = "";

    string q2Cut = "";
    if (UseOnlyMuMuEvents)      q2Cut = getMuMucut();
    else if (UseOnlyJpsiEvents) q2Cut = getJpsicut();
    else if (SplitInQ2){
        coutWarning("Using SplitInQ2! Always use with MuMu events!");
        if (UseLowQ2Range)      q2Cut = "(Q2 < 0.98e6 || (1.1e6 < Q2 && Q2 < 8.0e6))";    //q2Cuts = "nDiMuonMassBin < 9  || nDiMuonMassBin != 2";
        else                    q2Cut = "((11.0e6 < Q2 && Q2 < 12.5e6) || Q2 > 15.0e6)";    //q2Cuts = "nDiMuonMassBin > 9 || nDiMuonMassBin != 12";
    }
    else q2Cut = "(Q2 >0)"; //dummy cut in order to not having to deal with && at a beggining of the string
    finalCut.append(q2Cut);

    if (useExtraVar){
        string extraVarCut = "";
        double lowCut = extraVar.isEquidistant ? extraVar.binEdgesEquidistant.at(nExtraBin) : extraVar.binEdges.at(nExtraBin);
        double upperCut = extraVar.isEquidistant ? extraVar.binEdgesEquidistant.at(nExtraBin+1) : extraVar.binEdges.at(nExtraBin+1);
        extraVarCut = "(" + to_string(lowCut) + " < " + extraVar.cut+" && " + extraVar.cut+ "<= " + to_string(upperCut) + ")";
        coutDebug(extraVarCut);
        finalCut.append(" && ");
        finalCut.append(extraVarCut);
    }

    if (TMVAcut > -1.0){
        finalCut.append(" && ");
        finalCut.append(" (" + string(TMVAmethod) + "response > " + to_string(TMVAcut) + ")");
    }
    coutTest("Main cut: " + finalCut);
    return finalCut ;
}

string getAloneBranch(bool MC, bool TM, string customTMbranch, bool gammaTM){
    string customTag = "";
    if (MC){
        if (!TM) customTag = "NotTM";
        else {
            if (customTMbranch == "") customTMbranch  = TMbranch;
            if (customTMbranch == "TMedBKGCAT") customTag = "";
            if (customTMbranch == "TMed"){
                if (gammaTM) customTag = "_TMed";
                else         customTag = "_TMed_rndGamma";
            }
        }
    }
    return "IsAloneAt" + customTag;
}

string getMultipleCut(bool MC, bool TM, string customTMbranch, bool gammaTM, double TMVAcut){
    string aloneBranch = getAloneBranch(MC,TM,customTMbranch, gammaTM);
    string cut = (TMVAcut > -1.0) ?  string("<=" + to_string(TMVAcut)) : "==0.0";
    coutTest("Alone cut: " + aloneBranch + cut);
    return aloneBranch + cut;

}

string getFinalCut(bool MC, bool TM, bool notTM, string customTMbranch, bool gammaTM, bool UseOnlyMuMuEvents, bool UseOnlyJpsiEvents, bool SplitInQ2, bool UseLowQ2Range, bool useExtraVar, TMefficiencyClass extraVar,int nExtraBin, double TMVAcut, bool removeMultiple){

    string fullCut = getFullCut(UseOnlyMuMuEvents, UseOnlyJpsiEvents, SplitInQ2, UseLowQ2Range, useExtraVar, extraVar, nExtraBin, TMVAcut);
    string TMcut = MC ? getTMcut(TM,notTM,customTMbranch,gammaTM) : "";
    if (TMcut != "") TMcut = " && (" + TMcut + ")";
    string multipleCut = removeMultiple ? string( " && (" +getMultipleCut(MC,TM,customTMbranch,gammaTM,TMVAcut) + ")") : "";

    return fullCut + TMcut + multipleCut;
}

//Buggy 2015 MC needs to be removed and replaced by 2016
//TODO: make sure it's not loaded twice everywhere
string checkIf2015MC(string year, bool MC, bool Reference, bool PHSP){
    if (Kst2Kpluspi0Resolved && MC && !Reference && !PHSP && year== "2015") return "2016";
    else return year;
}

///////////////////////////////////////////////
//
//          Sanity checks
//
/////////////////////////////////////////////

bool checkMC(bool &MC, bool ReferenceChannel, bool PHSP, bool mutuallyExclusive){
    if (!MC && PHSP){//if PHSP or reference channel, set MC to true
        coutWarning("PHSP or Reference channel option selected, setting MC to true");
        MC = true;
    }
    if (mutuallyExclusive && ReferenceChannel && PHSP){
        coutERROR("You cannot select both Reference channel and PHSP! Abort.");
        return false;
    }
    return true;
}
bool checkMC(bool ReferenceChannel, bool PHSP){
    if (ReferenceChannel && PHSP){
        coutERROR("You cannot select both Reference channel and PHSP! Abort.");
        return false;
    }
    return true;
}

bool checkQ2Range(bool UseOnlyJpsi, bool UseOnlyMuMu){
    if (UseOnlyJpsi && UseOnlyMuMu){
        coutERROR("Make up your mind! Either Jpsi or mumu channel, it cannot be both! Program failed");
        return false;
    }
    return true;
}

bool checkTM(bool MC, bool &TM, bool &nonTM, bool Preselected){
    if ( (TM || nonTM) && !MC){
        coutWarning("Data is not TruthMatched! Setting TM and nonTM to false!");
        TM = false; //Keep nonTM as it is for the MC fit to get the shape
    }
    if ( (TM || nonTM) && !Preselected){
        coutWarning("Stripped MC is not TruthMatched! Setting TM to false!");
        TM = false;
        nonTM = false;
    }
    if (TM && nonTM){
        coutERROR("Cannot fit TM==1 and TM==0 data at the same time! Make up your mind. Abort.");
        return false;
    }
    return true;
}

bool checkKshort(bool &KshortDecaysInVelo){
    if (!Kst2Kspiplus && KshortDecaysInVelo){
        coutWarning("No LL/DD tracks in KplusPizero channel! Setting  KshortDecaysInVelo to false!");
        KshortDecaysInVelo = false;
    }
    return true;
}

bool checkRefYear(string year){
    if (Kst2Kspiplus && (year == "2015" || year =="2017" || year == "2018" )){
        coutERROR("Reference channel is only available for 2011, 2012 and 2016!");
        return false;
    }
    if (!Kst2Kspiplus && (year =="2017" || year == "2018" )){
        coutERROR("Reference channel is only available for 2011, 2012, 2015 and 2016!");
        return false;
    }
    return true;
}

bool checkEntries(TTree *tree){
    Int_t nEvents    = tree->GetEntries();
    if(nEvents == 0){
        coutERROR("Entries in data TTree " + string(tree->GetName()) + " are empty. Abort!");
        return false;
    }
    return true;
}

bool checkYear(int year, bool MC, bool ReferenceChannel, bool PHSP){
    bool flag = false;
    if (!MC){
        if (year == 2011 || year == 2012|| year == 2015 || year == 2016 || year == 2017 || year == 2018  ) flag=true;
        else flag=false;
    }
    else{
        if (KshortChannel){ //KS
            if (ReferenceChannel){
                if (year == 2011 || year == 2012||  year == 2016) flag=true;
                else flag=false;
            }
            else if (PHSP){
                if (year == 2011 || year == 2012|| year == 2015 || year == 2016 || year == 2017 || year == 2018  ) flag=true;
                else flag=false;
            }
            else{
                if (year == 2011 || year == 2012|| year == 2015 || year == 2016 || year == 2017 || year == 2018  ) flag=true;
                else flag=false;
            }
        }
        else{ //K+pi0
            if (ReferenceChannel){
                if (year == 2011 || year == 2012|| year == 2015 || year == 2016) flag=true;
                else flag=false;
            }
            else if (PHSP){
                if (year == 2011 || year == 2012|| year == 2015 || year == 2016 || year == 2017 || year == 2018) flag=true;
                else flag=false;
            }
            else{
                if (year == 2011 || year == 2012|| year == 2015 || year == 2016 || year == 2017 || year == 2018) flag=true;
                else flag=false;
            }
        }
    }
    if (!flag) coutERROR("Wrong year input! Your input is " + to_string(year)+ ". Abort.");
    return flag;
}

bool checkRun(int Run){
    if(Run != 1 && Run != 2 && Run != 12){
        coutERROR("Invalid Run number given: >> " + to_string( Run) + " << . Use only Run 1, Run 2 or Run 12! Exit now.");
        return 0;
    }
    return 1;
}

///////////////////////////////////////////////
//
//          Get correct names and tags
//
/////////////////////////////////////////////

string getTMtag(string customTMbranch){
    if (customTMbranch == "") customTMbranch  = TMbranch;
    if (customTMbranch == "TMedBKGCAT") return "_BKGCAT";
    if (customTMbranch == "TMed") return "_IDTM";
    if (customTMbranch == "TMed_noPi0") return "_noPi0_IDTM";

    coutERROR("Wrong customTM branch! Returning string 'wrong'");
    return "wrong";
}

string getTMtag(string customTMbranch, bool gammaTM){
    if (customTMbranch == "") customTMbranch  = TMbranch;
    if (customTMbranch == "TMedBKGCAT") return "_BKGCAT";
    if (customTMbranch == "TMed") return "_IDTM"+ string(gammaTM ? "" : "_rndGamma");;
    if (customTMbranch == "TMed_noPi0") return "_noPi0_IDTM"+ string(gammaTM ? "" : "_rndGamma");;

    coutERROR("Wrong customTM branch! Returning string 'wrong'");
    return "wrong";
}

string getWeightName(string customTMbranch, bool gammaTM){
    if (customTMbranch == "") customTMbranch  = TMbranch;
    if (customTMbranch == "TMedBKGCAT") return "weight2D_"+firstMCweight;
    if (customTMbranch == "TMed")   return "weight2D_"+firstMCweight + "_TM" + string(gammaTM ? "" : "_rndGamma");
    if (customTMbranch == "TMed_noPi0")  return "weight2D_"+firstMCweight + "_noPi0TM";
    coutERROR("Wrong customTMbranch name! Returning an empty string.");
    return "";
}

std::string getDataTypeTag(bool MC, bool Reference, bool PHSP, bool B0, bool K1, bool Inc){
    if (!MC) return "Data";
    else{
        if (PHSP) return "MC PHSP";
        else{
            if (B0) return Reference ? "MC B0toKstJpsi" : "MC B0toKstMuMu";
            else if (K1) return Reference ? "MC BtoK1Jpsi" : "MC BtoK1MuMu";
            else if (Inc) return Reference ? "MC BtoXJpsi" : "MC BtoXMuMu";
            else return Reference ? "Reference MC" : "Signal MC";
        }
    }
}

string getDataTypeTag(bool MC, bool Reference, bool PHSP){
    return getDataTypeTag(MC, Reference, PHSP, false, false, false);
}

string getYearRunTag(int Run, string year){
    if (Run == 0) return "Year " + year;
    else return "Run " + to_string(Run);
}
/////////////////////////////////////////////
//
// Bin edges for a similarly populated bins
//
/////////////////////////////////////////////
void getSimilarlyPopulatedBins(TH1D *histogram, int desiredBins, int nBins){

    //rebin (THIS ASSUMES THE ORIGINAL BIN WIDTH IS ONE!!!)
    int maxBinContent = histogram->GetEntries()/desiredBins;
    TAxis *axis = histogram->GetXaxis();
    Double_t xEdges[desiredBins+1];
    xEdges[0] = axis->GetBinLowEdge(1);

    int i = 0;
    for (int b = 1; b <desiredBins+1; b++){
        int content = 0;
        while (content < maxBinContent){
            content = content+histogram->GetBinContent(i);
            i++;
        }
        xEdges[b] = axis->GetBinLowEdge(i);
        i++;
    }

    xEdges[desiredBins] = axis->GetBinUpEdge(nBins);
    for (int b = 0; b<desiredBins+1; b++){
        cout << xEdges[b] << endl;
    }

    return;
}

void getSimilarlyPopulatedBinsVar(string var, int nBins, double lowEdge, double highEdge, int desiredBins, string year, string magnet, int Run, bool ReferenceChannel, bool PHSP){

    gStyle -> SetOptStat(0);
    gROOT->SetBatch(kTRUE);
    TH1::SetDefaultSumw2(kTRUE);

    //Load all files
    coutInfo("Opennig the trees");
    TChain *tree= new TChain("DecayTreeTruthMatched");
    //Open the file(s)
    if (Run==0){
        tree->Add(returnFileAddress(year, getRunID(year), magnet, true, false, true, ReferenceChannel, PHSP, false).c_str());
        coutDebug("Adding " + returnFileAddress(year, getRunID(year), magnet, true, false, true, ReferenceChannel, PHSP, false));
    }
    else{
        for (auto yr: yearsMC(ReferenceChannel, PHSP, Run)){
            tree->Add(returnFileAddress(yr, getRunID(yr), magnet, true, false, true, ReferenceChannel, PHSP, false).c_str());
            coutDebug("Adding " + returnFileAddress(yr, getRunID(yr), magnet, true, false, true, ReferenceChannel, PHSP, false));
        }
    }

    //activate all branches
    tree->SetBranchStatus("*",1);

    //Draw the histogram
    TH1D *h_tmp  = new TH1D ("h_tmp", "h_tmp", nBins,lowEdge,highEdge);
    tree->Draw(Form("%s>>h_tmp",var.c_str()));

    //Get the binnings
    getSimilarlyPopulatedBins(h_tmp,desiredBins, nBins);

    //free da memory
    delete h_tmp;
    delete tree;

    return;
}


/////////////////////////////////////////////
//
// Get result from a fit file
//  Comment out the content of the functions in order to compile getPathForPython *facepalm
//
/////////////////////////////////////////////

RooFitResult* getResult(TFile *fitFile){
    RooFitResult* fitResult;
    fitFile->GetObject("fitresult_pdf_data",fitResult);
    if (verboseLevel < 2) fitResult->Print();
    return fitResult;

}


RooRealVar* getVarFromResult(RooFitResult *fitResult, string name){
    RooRealVar *myVar = (RooRealVar*)fitResult->floatParsFinal().find(name.c_str());
    return myVar;
}

double getBplusMeanFromResult(TFile* fitFile){
    RooRealVar *mass = getVarFromResult(getResult(fitFile),"sig_mean");
    return mass->getVal();
}

//Read the yields from a fit file
double getSigYield(TFile *fitFile){
    TVectorD *yield = (TVectorD*)fitFile->Get("yield");
    return (*yield)[0];
}
double getSigYieldErr(TFile *fitFile){
    TVectorD *yield_err = (TVectorD*)fitFile->Get("yield_err");
    return (*yield_err)[0];
}
double getBkgYield(TFile *fitFile){
    TVectorD *bkg = (TVectorD*)fitFile->Get("background");
    return (*bkg)[0];
}
double getBkgYieldErr(TFile *fitFile){
    TVectorD *bkg_eff = (TVectorD*)fitFile->Get("background_err");
    return (*bkg_eff)[0];
}
double getEffSigma(TFile *fitFile){
    TVectorD *sigma = (TVectorD*)fitFile->Get("sigma_eff");
    return (*sigma)[0];
}

/////////////////////////////////////////////
//
// Histogram/graph helpers
//
/////////////////////////////////////////////

TH1D *convertTGraph(TGraph *graph){

    //Possibly fix tihs at some point to be generally repflecting TMefficiencyClass
    double binEdges[7] = {0.1e6, 4.0e6, 8.0e6, 11.0e6, 12.5e6, 15.0e6, 20.0e6};
    TH1D *hist = new TH1D("hist","hist",6,binEdges);

    auto nPoints = graph->GetN(); // number of points in your TGraph
    for(int i=0; i < nPoints; ++i) {
       double x,y;
       graph->GetPoint(i, x, y);
       hist->Fill(x,y); // ?
    }
    for (int b=0; b < hist->GetXaxis()->GetNbins(); b++){
        hist->SetBinError(b+1,0);
    }
    return hist;
}

/////////////////////////////////////////////
//
// Get numbers of MC events
//
/////////////////////////////////////////////
double generated_basic_events = 10000.0;
double generated_basic_events_PHSP = 1000.0;
double get_generated_events(bool PHSP){
    if (PHSP) return generated_basic_events_PHSP;
    else      return generated_basic_events;
}

vector<int> generated_signal_events_down    = {507551,  514015,  0,  1000281,  1151738, 1235528};
vector<int> generated_signal_events_up      = {502787,  500458,  0,  1000281, 1153816,	1196153};
vector<int> generated_reference_events_down = {1011831,	1003888, 1007712, 1010609, 0,       0};
vector<int> generated_reference_events_up   = {1007920, 1000278, 1009484, 1000281, 0,       0};
vector<int> generated_PHSP_events_down      = {85736,   226933,  95323,   264917,  246457,  291850};
vector<int> generated_PHSP_events_up        = {86004,   226430,  94980,   267674,  242673,  299252};

vector<double> gen_tables_signal_efficiency_down    = {0.14388, 0.14769, 0.1615, 0.1610, 0.1609, 0.1605};//RUN I taken from Reference Channel
vector<double> gen_tables_signal_efficiency_up      = {0.14422, 0.14787, 0.1608, 0.1611, 0.1595, 0.1609};//RUN I taken from Reference Channel
vector<double> gen_tables_reference_efficiency_down = {0.14388, 0.14769, 0.1581, 0.1585, 0.1585, 0.1585}; //2017+2018 taken as 2016
vector<double> gen_tables_reference_efficiency_up   = {0.14422, 0.14787, 0.1574, 0.1590, 0.1590, 0.1590}; //2017+2018 taken as 2016
vector<int> yield_signal_events        = {5451,  4906,  4096,  10382, 15994, 14919};
vector<int> yield_signal_events_err    = {70,    64,    1,     96,    15,    116};
vector<int> yield_reference_events     = {19730, 17579, 18392, 23965, 0,     0};
vector<int> yield_reference_events_err = {135,   127,   130,   149,   0,     0};
vector<int> yield_PHSP_events          = {14091, 33086, 13530, 48135, 54030, 61641};
vector<int> yield_PHSP_events_err      = {113,   174,   111,   209,   221,   237};

vector<int> yield_signal_events_notTM        = {7056,  6587,  6354,  15652, 20207, 19734};
vector<int> yield_signal_events_notTM_err    = {456,   61,    81,    55,    368,   222};
vector<int> yield_reference_events_notTM     = {24375, 22698, 23403, 30303, 0,     0};
vector<int> yield_reference_events_notTM_err = {236,   123,   143,   12,    0,     0};

vector<int> yield_signal_events_notTM_unique        = {5854,  5304,  5155,  13086, 17055, 16374};
vector<int> yield_signal_events_notTM_unique_err    = {89,    106,   86,    169,   172,   211};
vector<int> yield_reference_events_notTM_unique     = {20564, 18772, 19204, 25140, 0,     0};
vector<int> yield_reference_events_notTM_unique_err = {210,   123,   89,    167,   0,     0};

int get_position_from_year(string year){
    int position = year.back() - '0'; //get last char of a string and convert it to an int
    return position>4 ? position-3 : position-1; //11->0, 12->1, 15->2, 16->3, ...
}

int get_gen_evts(string year, bool ReferenceChannel, bool PHSP){
    int pos = get_position_from_year(year);
    if (ReferenceChannel) return generated_reference_events_down.at(pos) + generated_reference_events_up.at(pos);
    if (PHSP) return generated_PHSP_events_down.at(pos) + generated_PHSP_events_up.at(pos);
    return generated_signal_events_down.at(pos) + generated_signal_events_up.at(pos);
}

int get_gen_evts(int Run, bool ReferenceChannel, bool PHSP){
    int yield = 0;
    for (auto yr: yearsMC(ReferenceChannel,PHSP,Run)) yield += get_gen_evts(yr,ReferenceChannel,PHSP);
    return yield;
}

int get_selected_evts(string year, bool ReferenceChannel, bool PHSP){
    int pos = get_position_from_year(year);
    if (ReferenceChannel) return yield_reference_events.at(pos);
    if (PHSP) return yield_PHSP_events.at(pos);
    return yield_signal_events.at(pos);
}

int get_selected_evts(int Run, bool ReferenceChannel, bool PHSP){
    int yield = 0;
    for (auto yr: yearsMC(ReferenceChannel,PHSP,Run)) yield += get_selected_evts(yr,ReferenceChannel,PHSP);
    return yield;
}

int get_selected_evts_err(string year, bool ReferenceChannel, bool PHSP){
    int pos = get_position_from_year(year);
    if (ReferenceChannel) return yield_reference_events_err.at(pos);
    if (PHSP) return yield_PHSP_events_err.at(pos);
    return yield_signal_events_err.at(pos);
}

double get_selection_efficiency(string year, bool ReferenceChannel, bool PHSP){
    return double(get_selected_evts(year,ReferenceChannel,PHSP))/double(get_gen_evts(year,ReferenceChannel,PHSP));
}

double get_selection_efficiency(int Run, bool ReferenceChannel, bool PHSP){
    return double(get_selected_evts(Run,ReferenceChannel,PHSP))/double(get_gen_evts(Run,ReferenceChannel,PHSP));
}

double get_tables_eff(string year, bool ReferenceChannel){ //Get efficiency in a year as an average between up and down
    int pos = get_position_from_year(year);
    if (ReferenceChannel) return (gen_tables_reference_efficiency_down.at(pos)+gen_tables_reference_efficiency_up.at(pos))/2;
    else                  return (gen_tables_signal_efficiency_down.at(pos)+gen_tables_signal_efficiency_up.at(pos))/2;
}

double get_tables_eff(int Run, bool ReferenceChannel){ //Get efficiency of a run as an average between years
    double sum = 0;
    int nEntries = 0;
    double tmp_eff = 0;
    for (auto yr: yearsMC(ReferenceChannel,false,Run)){
        tmp_eff = get_tables_eff(yr, ReferenceChannel);
        if (tmp_eff!=0) nEntries++;
        sum += tmp_eff;
    }
    coutDebug("Total eff/nEntries: " + to_string(sum) + "/" + to_string(nEntries) );
    if (nEntries == 0) return 0;
    return sum/nEntries;
}

void print_all_yields_and_efficiencies(bool ReferenceChannel, bool PHSP){
    for (auto yr : yearsMC(ReferenceChannel,PHSP, 12)){
        cout << "------ year " << yr << " ------" << endl;
        cout << "\t\t Generated events:\t" << get_gen_evts(yr, ReferenceChannel, PHSP) << endl;
        cout << "\t\t Selected  events:\t" << get_selected_evts(yr, ReferenceChannel, PHSP) << endl;
        cout << "\t\t Efficiency: \t" << get_selection_efficiency(yr, ReferenceChannel, PHSP) << endl;
    }
    cout << "------ Run 1 ------" << endl;
    cout << "\t\t Generated events:\t" << get_gen_evts(1, ReferenceChannel, PHSP) << endl;
    cout << "\t\t Selected  events:\t" << get_selected_evts(1, ReferenceChannel, PHSP) << endl;
    cout << "\t\t Efficiency: \t" << get_selection_efficiency(1, ReferenceChannel, PHSP) << endl;
    cout << "------ Run 2 ------" << endl;
    cout << "\t\t Generated events:\t" << get_gen_evts(2, ReferenceChannel, PHSP) << endl;
    cout << "\t\t Selected  events:\t" << get_selected_evts(2, ReferenceChannel, PHSP) << endl;
    cout << "\t\t Efficiency: \t" << get_selection_efficiency(2, ReferenceChannel, PHSP) << endl;
    return;
}

//----------------------------------------------------------------------------------------------------------

#endif // UTILS_HPP