//Class containing value and errors used for efficienceis
//Renata Kopecna

#include "GlobalFunctions.hh"
#include "Paths.hpp"
#include "Design.hpp"
#include "MassFit.cpp"

using namespace std;
using namespace RooFit;
using namespace RooStats;

//////////////////////////////////////////////////////
///  class EffandError()
//////////////////////////////////////////////////////
///
/// Used for storing efficiency + hiighError - lowError
///
///
/////////////////////////////////////////////////////////
///  GetEfficiency
//////////////////////////////////////////////////////
///
/// Reads two given files, reads the yields (TVectorDouble) from them
///
///
/////////////////////////////////////////////////////////
///  getTruthMatchingEfficiencySimple
//////////////////////////////////////////////////////
///
/// Returns EffAndError for Truthmatching/selected,
/// either per year or per run
///
///
/////////////////////////////////////////////////////////
///  getBDTEfficiencySimple
//////////////////////////////////////////////////////
///
/// Returns EffAndError for a given TMVAresponse cut,
/// done not by fitting the mass peak but by a simple count,
/// either per year or per run
///
///
/////////////////////////////////////////////////////////
///  getSelectionEfficiencySimple
//////////////////////////////////////////////////////
///
/// Returns EffAndError for a given year and Run
/// from generated and selected TM events
///
///


//Define EffAndError class
//----------------------------------------------------------------------------------------------------------
class EffAndError{
 public:
    double value;
    double highError;
    double lowError;
    EffAndError(){ //default constructor
        value = 0;
        highError = 0;
        lowError = 0;
    }
    EffAndError(double val){ //default constructor
        value = val;
        highError = 0;
        lowError = 0;
    }
    EffAndError(double val,double highErr, double lowErr){ //default constructor
        value = val;
        highError = highErr;
        lowError = lowErr;
    }
    ~EffAndError();   //destuctor
};

EffAndError::~EffAndError(){//destuctor

}
string mainSignalShape = "OneCB";


//Calculate efficiency from two files
//----------------------------------------------------------------------------------------------------------
EffAndError GetEfficiency(TFile *fitFileNumerator, string nameNumerator, TFile *fitFileDenominator, string nameDenominator){

    if (fitFileNumerator == NULL){
        coutERROR("File " + string(fitFileNumerator->GetPath()) + " not opened!");
        return EffAndError(0,0,0);
    }

    if (fitFileDenominator == NULL){
        coutERROR("File " + string(fitFileDenominator->GetPath()) + " not opened!");
        return EffAndError(0,0,0);
    }

    // Get yield for numerator data
    TVectorD *SigYieldNumeratorVec = (TVectorD*)fitFileNumerator->Get("yield");
    double SigYieldNumerator  = (*SigYieldNumeratorVec)[0];

    // Get yield for denominator data
    TVectorD *SigYieldDenominatorVec = (TVectorD*)fitFileDenominator->Get("yield");
    double SigYieldDenominator  = (*SigYieldDenominatorVec)[0];

    double efficiency = SigYieldNumerator/SigYieldDenominator;

    coutDebug("Yield from " + nameNumerator + " fit:\t" + to_string(SigYieldNumerator));
    coutDebug("Yield from " + nameDenominator + " fit:\t" + to_string(SigYieldDenominator));
    coutInfo("Efficiency:\t" + to_string(efficiency));

    //EffAndError result = EffAndError(efficiency,EffErrHi,EffErrLo);
    EffAndError result = EffAndError(efficiency,0.0,0.0);

    coutDebug("Efficiency succesfully calculated" );

    return result;

}

//Get efficiency from counts in MC (TM is possible since background in non-TMed sample is hard to model anyway)
//----------------------------------------------------------------------------------------------------------

EffAndError getTMEfficiencySimple(string year, int Run, bool UseOnlyMuMuEvents, bool PHSP, bool KshortDecaysInVelo,  bool RemoveMultiple, bool weighted,  string sExtraVar, int nExtraBin, string customTMbranch, bool gammaTM){  //TODO: add customTMbranch and gammaTM

    gROOT->SetBatch(); //ROOT stops plotting canvases
    EffAndError eff = EffAndError(1.0,0.0,0.0);
    bool useRefChannel = !PHSP && !UseOnlyMuMuEvents;
    //check if using reference channel for not available years
    if (!checkYear(std::stoi(year),true,useRefChannel,PHSP)) return eff;

    //load the tree from the file
    TChain * T = new TChain("DecayTreeTruthMatched");

    std::vector<string> years;
    if (Run == 0)years.push_back(year); //If Run==0, calculate efficiecny for separate year
    else years = yearsMC(useRefChannel,PHSP,Run); //Otherwise combine the years together

    for (auto& yr : years){
        T->Add(GetBDToutputFile(yr,getRunID(yr),true,useRefChannel,PHSP,KshortDecaysInVelo,false,false).c_str());
        coutDebug("Adding " + GetBDToutputFile(yr,getRunID(yr),true,useRefChannel,PHSP,KshortDecaysInVelo,false,false));
    }

    bool useExtraVar = useExtraVarBool(sExtraVar); //perform once so it doesn't have to run through all the ifs every time

    TMefficiencyClass extraVar = useExtraVar ? TMefficiencyClass(sExtraVar) : TMefficiencyClass();

    // Set cuts on the tree
    string massCut = "";
    if (UseDTF) massCut = " ( B_plus_M_DTF > " + to_string(PDGMASS.B_PLUS-100) +" &&  B_plus_M_DTF <" + to_string(PDGMASS.B_PLUS+100)+ ") ";
    else massCut =  " ( B_plus_M > " + to_string(PDGMASS.B_PLUS-100) +" &&  B_plus_M <" + to_string(PDGMASS.B_PLUS+100)+ ") ";
    coutDebug("Mass cut: " +massCut);

    string sCut = getFinalCut(true, false, false, customTMbranch, gammaTM, UseOnlyMuMuEvents, useRefChannel, false, false, useExtraVar, extraVar, nExtraBin, -1, RemoveMultiple);
    coutDebug("sCut: " + sCut);
    sCut.append(" && ");
    sCut.append(massCut);
    if (weighted) sCut = getWeightName(customTMbranch,gammaTM)+"*(" + sCut + ")";
    coutDebug("First cut: " + sCut);

    //Get number of entries before TM
    TH1F *h_noTM = new TH1F("h_noTM","h_noTM",1002,-1.0,1.0);
    T->Draw(TMVAmethod+"response >> h_noTM", sCut.c_str());

    sCut = getFinalCut(true, true, false, customTMbranch, gammaTM, UseOnlyMuMuEvents, useRefChannel, false, false, useExtraVar, extraVar, nExtraBin, -1, RemoveMultiple);
    sCut.append(" && ");
    sCut.append(massCut);
    if (weighted) sCut = getWeightName(customTMbranch,gammaTM)+ "*(" + sCut + ")";
    coutDebug("Second cut: " + sCut);

    //Get number of entries after the BDT cut
    TH1F *h_TM = new TH1F("h_TM","h_TM",1002,-1.0,1.0);
    T->Draw(TMVAmethod+"response >>h_TM", sCut.c_str());

    //Calculate efficiency and its binomial error
    double passEntries = h_TM->Integral();
    double nentries = h_noTM->Integral();
    double effErrHi = TMath::Sqrt(passEntries*(1.0-passEntries/nentries))/nentries;
    eff.value = passEntries/ nentries;

    coutDebug("All MC entries: " + to_string(nentries));
    coutDebug("MC entries passing TM: " + to_string(passEntries));
    coutDebug("TM efficiency: " + to_string(eff.value));

    eff.highError = effErrHi;
    eff.lowError =  effErrHi;

    h_noTM->Clear();
    h_TM->Clear();
    delete T;

    return eff;

}

EffAndError getBDTEfficiencySimple(string year, Double_t TMVAcut, int Run, bool UseOnlyMuMuEvents,bool PHSP,
                     bool KshortDecaysInVelo, bool IncludeMultipleEff, bool weighted, string sExtraVar, int nExtraBin){ //TODO: check weight

    EffAndError eff = EffAndError(1.0,0.0,0.0);

    bool useRefChannel = !PHSP && !UseOnlyMuMuEvents;

    //check if using reference channel for not available years
    if (!checkYear(std::stoi(year),true,useRefChannel,PHSP)) return eff;

    //Get low TMVA response lower boundary
    double lowBDTcut = -1.0; //check if at lower limit
    if (TMVAmethod == "MLP") lowBDTcut = 0.0;
    if (TMVAcut == lowBDTcut){
        coutERROR("BDT cut is 100% effective at "+ to_string(lowBDTcut) + "! returning efficiency 1+-0.");
        return eff;
    }

    //load the tree from the file
    TChain * T = new TChain("DecayTreeTruthMatched");

    std::vector<string> years;
    if (Run == 0){ //If Run==0, calculate efficiecny for separate year
        years.push_back(year);
    }
    else years = yearsMC(useRefChannel,PHSP,Run); //Otherwise combine the years together

    for (auto& yr : years){
        int RunID = (yr == "2011" || yr == "2012") ? 1:2;
        T->Add(GetBDToutputFile(yr,RunID,true,useRefChannel,PHSP,KshortDecaysInVelo,false,false).c_str());
    }

    bool useExtraVar = useExtraVarBool(sExtraVar); //perform once so it doesn't have to run through all the ifs every time
    TMefficiencyClass extraVar = useExtraVar ? TMefficiencyClass(sExtraVar) : TMefficiencyClass();

    gROOT->SetBatch(); //ROOT stops plotting canvases
    // Set cuts on the tree
    TString massCut = "";
    if (UseDTF) massCut = " ( B_plus_M_DTF > " + to_string(PDGMASS.B_PLUS-100) +" &&  B_plus_M_DTF <" + to_string(PDGMASS.B_PLUS+100)+ ") ";
    else massCut =  " ( B_plus_M > " + to_string(PDGMASS.B_PLUS-100) +" &&  B_plus_M <" + to_string(PDGMASS.B_PLUS+100)+ ") ";

    //add vetos
    string sCut = getFinalCut(true, true, false, "", gammaTMdefault, UseOnlyMuMuEvents, useRefChannel, false, false, useExtraVar, extraVar, nExtraBin, lowBDTcut, !IncludeMultipleEff);
    string LLorDD = string("&& (KshortDecayInVeLo == ") + (KshortDecaysInVelo ? "1":"0") + ")";
    sCut.append(" && ");
    sCut.append(massCut);
    if(Kst2Kspiplus) sCut.append(LLorDD);
    if (weighted) sCut = getWeightName("",gammaTMdefault)+"*(" + sCut + ")";
    coutDebug("Using cut: "+sCut);
    //TODO!!!!! Naming convension for including multiple efficiency!!!!!

    //Get number of entries before the BDT cut
    TH1F *h_BDTresponseMCAll = new TH1F("h_"+TMVAmethod+"responseAll","h_"+TMVAmethod+"responseAll",1002,lowBDTcut,1.0);
    T->Draw(TMVAmethod+"response >> h_"+TMVAmethod+"responseAll", sCut.c_str());

    sCut = getFinalCut(true, true, false, "", gammaTMdefault, UseOnlyMuMuEvents, useRefChannel, false, false, useExtraVar, extraVar, nExtraBin, TMVAcut, true);
    sCut.append(" && ");
    sCut.append(massCut);
    if (Kst2Kspiplus) sCut.append(LLorDD);
    if (weighted) sCut = getWeightName("",gammaTMdefault)+"*(" + sCut + ")";
    coutDebug("Using cut: "+sCut);

    //Get number of entries after the BDT cut
    TH1F *h_BDTresponseMCtmp = new TH1F("h_"+TMVAmethod+"response","h_"+TMVAmethod+"response",1002,lowBDTcut,1.0);
    T->Draw(TMVAmethod+"response >> h_"+TMVAmethod+"response", sCut.c_str());

    //Calculate efficiency and its binomial error
    double passEntries = h_BDTresponseMCtmp->Integral();
    double nentries = h_BDTresponseMCAll->Integral();
    double effErrHi = TMath::Sqrt(passEntries*(1.0-passEntries/nentries))/nentries;
    eff.value = passEntries/ nentries;

    coutDebug("All MC entries: " + to_string(nentries));
    coutDebug("MC entries passing TMVA cut " + to_string(TMVAcut) + ": " + to_string(passEntries));
    coutDebug("TMVA cut" + to_string(TMVAcut) + " efficiency: " + to_string(eff.value));

    eff.highError = effErrHi;
    eff.lowError =  effErrHi;

    h_BDTresponseMCAll->Clear();
    h_BDTresponseMCtmp->Clear();

    delete T;

    coutDebug("BDT efficiency for cut at " + to_string(TMVAcut) + " calculated!" );
    return eff;

}

//Get selection efficiency from counts in MC
//----------------------------------------------------------------------------------------------------------

EffAndError getSelectionEfficiencySimple(bool full, string year, int Run, bool  UseOnlyMuMuEvents, bool useRefChannel, bool PHSP, bool KshortDecaysInVelo,  bool RemoveMultiple, bool weighted, string sExtraVar, int nExtraBin, string customTMbranch, bool gammaTM){  //For studies on variables

    gROOT->SetBatch(); //ROOT stops plotting canvases
    EffAndError eff = EffAndError(1.0,0.0,0.0);
    //check if using reference channel for not available years
    if (!checkYear(std::stoi(year),true,useRefChannel,PHSP)) return eff;

    //Scaling factor, pretty much the number of generated events
    //For easier division it's a double right away
    double scale = get_generated_events(PHSP);

    std::vector<string> years;
    if (Run == 0)years.push_back(year); //If Run==0, calculate efficiecny for separate year
    else years = yearsMC(useRefChannel,PHSP,Run); //Otherwise combine the years together

    //load the tree from the file
    TChain * T = new TChain("DecayTreeTruthMatched");
    for (auto& yr : years){
        T->Add(GetBDToutputFile(yr,getRunID(yr),true,useRefChannel,PHSP,KshortDecaysInVelo,false,false).c_str());
        coutDebug("Adding " + GetBDToutputFile(yr,getRunID(yr),true,useRefChannel,PHSP,KshortDecaysInVelo,false,false));
    }

    //Kst2Kpluspi0Resolved case
    TChain* OG = 0;
    if(Kst2Kpluspi0Resolved){
        OG=new TChain("DecayTree");
        OG->Add(GetGenLevelFile(useRefChannel,PHSP).c_str());
        coutDebug("Adding " + GetGenLevelFile(useRefChannel,PHSP));
    }
    else{
        coutERROR("Not implemented yet.");
        return EffAndError();
    }

    bool useExtraVar = useExtraVarBool(sExtraVar); //perform once so it doesn't have to run through all the ifs every time
    TMefficiencyClass extraVar = useExtraVar ? TMefficiencyClass(sExtraVar) : TMefficiencyClass();

    // Set cuts on the tree
    string massCut = "";
    if (UseDTF) massCut = " ( B_plus_M_DTF > " + to_string(PDGMASS.B_PLUS-100) +" &&  B_plus_M_DTF <" + to_string(PDGMASS.B_PLUS+100)+ ") ";
    else massCut =  " ( B_plus_M > " + to_string(PDGMASS.B_PLUS-100) +" &&  B_plus_M <" + to_string(PDGMASS.B_PLUS+100)+ ") ";
    coutDebug("Mass cut: " +massCut);

    string sCut = getFinalCut(true, true, false, customTMbranch, gammaTM, UseOnlyMuMuEvents, useRefChannel, false, false, useExtraVar, extraVar, nExtraBin, -1, RemoveMultiple);
    sCut.append(" && ");
    sCut.append(massCut);
    if (weighted) sCut = getWeightName(customTMbranch,gammaTM)+ "*(" + sCut + ")";
    coutDebug("Full cut: " + sCut);

    string minCut = getFinalCut(true,false,false,"",false,UseOnlyMuMuEvents,useRefChannel,false,false, useExtraVar, extraVar, nExtraBin, -1.0, false);
    //Get number of entries before anything
    coutDebug("Minimal cut: " + minCut);

    TH1F *h_OG = new TH1F("h_OG","h_OG",5,5270.0,5290.0); //The B_plus mass is at PDG value in the sample
    //coutDebug("Entries in h_OG: "+to_string(OG->GetEntries()));
    OG->Draw("B_plus_M >> h_OG", minCut.c_str());


    //Get number of entries after the BDT cut
    TH1F *h_TM = new TH1F("h_TM","h_TM",1002,-1.0,1.0);;
    T->Draw(TMVAmethod+"response >>h_TM", sCut.c_str());

    //Calculate efficiency and its binomial error
    double passEntries = h_TM->Integral();
    double nentries = h_OG->Integral();
    int totalEntries = Run == 0? get_gen_evts(year,useRefChannel,PHSP) : get_gen_evts(Run,useRefChannel,PHSP);
    nentries = nentries*totalEntries/scale;

    double effErrHi = TMath::Sqrt(passEntries*(1.0-passEntries/nentries))/nentries; //TODO error calculation
    eff.value = passEntries/ nentries;

    //Multiply by generator-level efficiency if requested
    if (full) {
        double genEvtEff = (Run==0) ? get_tables_eff(year,useRefChannel) : get_tables_eff(Run,useRefChannel);
        if (genEvtEff==0.0) genEvtEff = 1.0;
        eff.value = eff.value*genEvtEff;
        effErrHi = effErrHi*genEvtEff;
    }

    coutDebug("All EvtGen entries: " + to_string(h_OG->Integral()));;
    coutDebug("All EvtGen scaled entries: " + to_string(nentries));
    coutDebug("MC entries passing TM: " + to_string(passEntries));
    coutDebug("Selection efficiency: " + to_string(eff.value));


    eff.highError = effErrHi;
    eff.lowError =  effErrHi;

    h_OG->Clear();
    h_TM->Clear();
    delete T;
    delete OG;

    return eff;

}