EWP-BplusToKstMuMu-AngAna/Code/Selection/Efficiency.cpp

//Functions used to obtain various efficiencies in the K+pi0 channel
//Renata Kopecna

#include "GlobalFunctions.hh"
#include "EfficiencyClass.cpp"
#include "Design.hpp"
#include "Paths.hpp"
#include "Utils.hpp"

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

double defaultBDTstep = 0.005;

// Small helper function
//----------------------------------------------------------------------------------------------------------
Double_t getLowBDTcut(bool fineScan){
    if (fineScan) return 0.9;
    else return (TMVAmethod == "MLP") ? 0.0 : -1.0;
}
TH1D *convertTGraph(TGraphErrors *effGraph){

    double *x, *y;
    double *ex, *ey;

    x = effGraph->GetX();
    y = effGraph->GetY();
    ex = effGraph->GetEX();
    ey = effGraph->GetEY();
    int N = effGraph->GetN();

    double binEdges[N+1];
    for (int i = 0; i<N; i++){
        binEdges[i] = x[i]-ex[i];
        coutDebug(to_string(binEdges[i]));
    }
    binEdges[N] = x[N-1]+ex[N-1];
    coutDebug(to_string(binEdges[N]));

    TH1D *hist = new TH1D("hist","hist",N,binEdges);
    for (int b = 0; b < N; b++){
        hist->SetBinContent(b+1,y[b]);
        hist->SetBinError(b,ey[b]);
    }

    hist->Draw();
    return hist;

}

//Get L0Muon efficiency
//----------------------------------------------------------------------------------------------------------
TH1D * getL0MuonEfficiency(string year, string magnet, bool MC, bool ReferenceChannel, bool PHSP, bool useDimuon){

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

    if (PHSP) MC = true;
    if (ReferenceChannel) MC = true;
    //Load all files
    TChain* tree = new TChain("DecayTree");

    if (magnet == "both"){
        string input_path = GetInputFile(year,"down",true,MC,ReferenceChannel,PHSP,false);
        replace(input_path,".root","_L0Muon.root");
        tree->Add(input_path.c_str());
        replace(input_path,"down","up");
        tree->Add(input_path.c_str());
    }
    else{
        string input_path = GetInputFile(year,magnet,true,MC,ReferenceChannel,PHSP,false);
        replace(input_path,".root","_L0Muon.root");
        tree->Add(input_path.c_str());
    }

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

    //Draw the histograms
    Double_t xEdges[16] = {0.0,1950.0,2254.25,2522.0,
                           2787.75,3061.75,3346.0,3658.75,
                           4003.75,4401.75,4868.5,5444.0,
                           6190.25,7275.25,9257.25,20000.0};

    TH1D *mu_max_all = new TH1D ("mu_max_all","mu_max_all",15,xEdges);
    TH1D *mu_max_pass = new TH1D ("mu_max_pass","mu_max_pass",15,xEdges);

    //In RunI nSPDHits < 600
    string cut = "max(mu_plus_PT,mu_minus_PT)< 20000 && (mu_plus_ProbNNmu > 0.25 && mu_minus_ProbNNmu > 0.25) && nSPDHits < 450";
    if (year == "2011" || year == "2012") replace(cut,"450","600");
    tree->Draw("max(mu_plus_PT,mu_minus_PT)>>mu_max_all",cut.c_str());

    if (useDimuon) cut = cut + " && (B_plus_L0MuonDecision_TOS==1 || B_plus_L0DiMuonDecision_TOS==1 )";
    else cut = cut + "&& B_plus_L0MuonDecision_TOS==1";

    tree->Draw("max(mu_plus_PT,mu_minus_PT)>>mu_max_pass",cut.c_str());

    //Add true muon pt possibly for MC
    mu_max_all = (TH1D*)gDirectory->Get("mu_max_all");
    mu_max_pass = (TH1D*)gDirectory->Get("mu_max_pass");

    string sMC = ( ReferenceChannel ? "RefMC" : (PHSP ? "PHSP" : "MC") );
    string L0MuonEffName = "L0MuonEff" + (MC ? sMC : "");
    TH1D *L0MuonEff = (TH1D*)mu_max_all->Clone(L0MuonEffName.c_str());
    L0MuonEff->Divide(mu_max_pass,mu_max_all);

    delete mu_max_all;
    delete mu_max_pass;

    return L0MuonEff;
}
int getAllL0MuonEfficiencies(string year, string magnet, bool useDimuon){

    TH1D *L0MuonEffData = getL0MuonEfficiency(year,magnet,false, false, false,useDimuon);
    TH1D *L0MuonEffMC = getL0MuonEfficiency(year,magnet,true, false, false,useDimuon);
    TH1D *L0MuonEffRefMC = getL0MuonEfficiency(year,magnet,false, true, false,useDimuon);
    TH1D *L0MuonEffPHSP = getL0MuonEfficiency(year,magnet,false, false, true,useDimuon);

    TH1D *L0MuonEffRatioMC =    (TH1D*)L0MuonEffData->Clone("L0MuonEffRatioMC");
    TH1D *L0MuonEffRatioRefMC = (TH1D*)L0MuonEffData->Clone("L0MuonEffRatioRefMC");
    TH1D *L0MuonEffRatioPHSP =  (TH1D*)L0MuonEffData->Clone("L0MuonEffRatioPHSP");
    L0MuonEffRatioMC->Divide(L0MuonEffMC);
    L0MuonEffRatioRefMC->Divide(L0MuonEffRefMC);
    L0MuonEffRatioPHSP->Divide(L0MuonEffPHSP);


    string output_path = GetEfficiencyFile("L0Muon",year,magnet,false,false,false,false,false,false,false,false,useDimuon ? "DiMuon" : string(""));
    TFile *output = new TFile(output_path.c_str(),"RECREATE");
    output->cd();

    TCanvas *c_plot = designL0MuonEfficiencies(year, magnet, L0MuonEffData,L0MuonEffPHSP,L0MuonEffRatioPHSP,useDimuon ? "DiMuon_PHSP" :"PHSP");
    TCanvas *c_plotMC = designL0MuonEfficiencies(year, magnet, L0MuonEffData,L0MuonEffMC,L0MuonEffRatioMC,useDimuon ? "DiMuon_MC" :"MC");
    TCanvas *c_plotRefMC = designL0MuonEfficiencies(year, magnet, L0MuonEffData,L0MuonEffRefMC,L0MuonEffRatioRefMC,useDimuon ? "DiMuon_RefMC" :"RefMC");

    L0MuonEffData->Write();
    L0MuonEffMC->Write();
    L0MuonEffRefMC->Write();
    L0MuonEffPHSP->Write();

    L0MuonEffRatioMC->Write();
    L0MuonEffRatioRefMC->Write();
    L0MuonEffRatioPHSP->Write();

    c_plot->Write();
    c_plotMC->Write();
    c_plotRefMC->Write();

    delete L0MuonEffData;
    delete L0MuonEffMC;
    delete L0MuonEffRefMC;
    delete L0MuonEffPHSP;

    delete L0MuonEffRatioMC;
    delete L0MuonEffRatioRefMC;
    delete L0MuonEffRatioPHSP;

    delete c_plot;
    delete c_plotMC;
    delete c_plotRefMC;


    output->Close();

    return 1;
}
int getAllL0MuonEfficienciesAllYears(int Run, string magnet, bool useDimuon){
    for (auto& year : yearsMC(false,true,Run)){
        if (getAllL0MuonEfficiencies(year,magnet,useDimuon)==0) return 0;
    }
    return 1;
}
//----------------------------------------------------------------------------------------------------------

//----------------------------------------------------------------------------------------------------------
//BDT scans
//----------------------------------------------------------------------------------------------------------
int ScanBDTEfficiency(string year, string magnet, Double_t BDTstep,
                     int Run, bool UseOnlyMuMuEvents, bool PHSP,
                     bool KshortDecaysInVelo, bool IncludeMultipleEff,
                     bool weighted){

    //Get low TMVA response lower boundary
    Double_t lowBDTcut = getLowBDTcut(false);

    TGraphErrors *effGraph = new TGraphErrors();
    gROOT->SetBatch(); //ROOT stops plotting canvases

    //loop over BDT cuts
    for(double fBDTcut = lowBDTcut+BDTstep; fBDTcut < 1.0; fBDTcut += BDTstep){
        EffAndError efficiency = EffAndError(); //put zeroes everywhere
        efficiency  = getBDTEfficiencySimple(year,fBDTcut,Run,UseOnlyMuMuEvents,PHSP,KshortDecaysInVelo,IncludeMultipleEff,weighted,"",-1);
        //Save the efficiency
        effGraph->SetPoint(effGraph->GetN(), fBDTcut, efficiency.value);
        effGraph->SetPointError(effGraph->GetN()-1,0,efficiency.highError);
        coutDebug("---"+string(TMVAmethod)+" cut "+to_string(fBDTcut)+" efficiency: "+to_string(efficiency.value)+"+-"+to_string(efficiency.highError));
    }

    //Write to file and close
    string path =  GetEfficiencyFile("BDT",year,magnet,Run,!UseOnlyMuMuEvents&&!PHSP, UseOnlyMuMuEvents,PHSP, KshortDecaysInVelo, IncludeMultipleEff,weighted, false, false,"");
    TFile * effFile = new TFile(path.c_str(), "RECREATE");
    coutInfo("Writting into a file " + path);
    effFile->cd();
    string effName ="";
    if (IncludeMultipleEff) effName = "{"+string(TMVAmethod)+"}^{Removed multiple}";
    else                    effName = "{"+string(TMVAmethod)+"}";

    string title ="";
    if (Run ==0) title = string(TMVAmethod)+"effScan_"+year;
    else         title = string(TMVAmethod)+"effScan_Run" + to_string(Run);
    designBasicEfficiency(effGraph,title,string(TMVAmethod)+" response",effName,0.5,1.025);
    effGraph->Write("",TObject::kWriteDelete);

    TCanvas * c1 = new TCanvas("c1", "c1");
    c1->cd();
    effGraph->Draw("ap");    
    addYearTag(0.25,0.55,year, Run, 1, 1.3);
    addAnyTag(0.25,0.45,getDataTypeTag(true,!UseOnlyMuMuEvents&&!PHSP,PHSP) ,1,1.3);
    replace(path,".root",".eps");
    c1->SaveAs(path.c_str(),"eps");
    effFile->Close();
    return 1;

}

int ScanBDTEfficiencyAllYears(bool UseOnlyMuMuEvents, bool PHSP, bool weighted,bool IncludeMultipleEff){
       for (auto yr: yearsMC(!UseOnlyMuMuEvents&&!PHSP,PHSP,12)){
           if (ScanBDTEfficiency(yr,"both",defaultBDTstep,0,UseOnlyMuMuEvents,PHSP,false,IncludeMultipleEff,weighted)==0) return 0;
           if (KshortChannel)
               if (ScanBDTEfficiency(yr,"both",defaultBDTstep,0,UseOnlyMuMuEvents,PHSP,true,IncludeMultipleEff,weighted)==0) return 0;
       }
       return 1;
}
int ScanMultipleCandidatesEfficiencyAllYearsAllSamples(bool weighted,bool IncludeMultipleEff){
    if (ScanBDTEfficiencyAllYears(true,  false, weighted, IncludeMultipleEff)==0) return 0;
    if (ScanBDTEfficiencyAllYears(false, false, weighted, IncludeMultipleEff)==0) return 0;
    if (ScanBDTEfficiencyAllYears(false, true,  weighted, IncludeMultipleEff)==0) return 0;
    return 1;
}

int ScanBDTEfficiencyAllSig (int Run, Double_t BDTstep, bool SplitByYears, bool RemoveMultiple, bool weighted){

    bool KshortDecaysInVelo = false;

    if(Run != 1 && Run != 2 && Run != 12){
        std::cout << "[ERROR]\t\tInvalid Run ID given: " << Run << ". Exit programm!" << std::endl;
        return 0;
    }
    if (SplitByYears){
        for (auto& year : yearsMC(false,false,Run))
            if (ScanBDTEfficiency(year,"both",BDTstep,0,true,false, KshortDecaysInVelo, RemoveMultiple, weighted)==0) return 0;
    }
    else{
        if (Run == 1 || Run == 12)  if (ScanBDTEfficiency("2011","both",BDTstep,1,true, false, KshortDecaysInVelo, RemoveMultiple, weighted)==0) return 0;
        if (Run == 2 || Run == 12)  if (ScanBDTEfficiency("2016","both",BDTstep,2,true, false, KshortDecaysInVelo, RemoveMultiple, weighted)==0) return 0;
    }

    return 1;
}
int ScanBDTEfficiencyAllPHSP(int Run, Double_t BDTstep, bool SplitByYears, bool RemoveMultiple, bool weighted){

    bool KshortDecaysInVelo = false;

    if(Run != 1 && Run != 2 && Run != 12){
        std::cout << "[ERROR]\t\tInvalid Run ID given: " << Run << ". Exit programm!" << std::endl;
        return 0;
    }
    if (SplitByYears){
        for (auto& year : yearsMC(false,true,Run))
            if (ScanBDTEfficiency(year,"both",BDTstep,0,true,true, KshortDecaysInVelo, RemoveMultiple, weighted)==0) return 0; //TODO: check if PHSP efficiency should be for mumu or overall!
    }
    else{
        if (ScanBDTEfficiency("2011","both",BDTstep,1,true, true, KshortDecaysInVelo, RemoveMultiple, weighted)==0) return 0;
        if (ScanBDTEfficiency("2016","both",BDTstep,2,true, true, KshortDecaysInVelo, RemoveMultiple, weighted)==0) return 0;
    }

    return 1;
}
int ScanBDTEfficiencyAllRef (int Run, Double_t BDTstep, bool SplitByYears, bool RemoveMultiple, bool weighted){

    bool KshortDecaysInVelo = false;

    if (SplitByYears){
        for (auto& year : yearsMC(true,false,Run))
            if (ScanBDTEfficiency(year,"both",BDTstep,0,false, false, KshortDecaysInVelo, RemoveMultiple, weighted)==0) return 0;
    }
    else{
        if (ScanBDTEfficiency("2011","both",BDTstep,1,false, false, KshortDecaysInVelo, RemoveMultiple, weighted)==0) return 0;
        if (ScanBDTEfficiency("2016","both",BDTstep,2,false, false, KshortDecaysInVelo, RemoveMultiple, weighted)==0) return 0;
    }

    return 1;
}

int ScanBDTEfficiencyAll(){
    ScanBDTEfficiencyAllSig(1,0.001,true,true,true);
    ScanBDTEfficiencyAllSig(2,0.001,true,true,true);
    ScanBDTEfficiencyAllSig(1,0.001,false,true,true);
    ScanBDTEfficiencyAllSig(2,0.001,false,true,true);

    ScanBDTEfficiencyAllRef(1,0.001,true,true,true);
    ScanBDTEfficiencyAllRef(2,0.001,true,true,true);
    ScanBDTEfficiencyAllRef(1,0.001,false,true,true);
    ScanBDTEfficiencyAllRef(2,0.001,false,true,true);

    ScanBDTEfficiencyAllPHSP(1,0.001,true,true,true);
    ScanBDTEfficiencyAllPHSP(2,0.001,true,true,true);
    ScanBDTEfficiencyAllPHSP(1,0.001,false,true,true);
    ScanBDTEfficiencyAllPHSP(2,0.001,false,true,true);

    return 1;
}

//Plot BDT efficiency at certain TMVA cut in a variable
//----------------------------------------------------
int plotBDTEfficiency(string year, int Run, Double_t TMVAcut, bool UseOnlyMuMuEvents, bool PHSP,  bool KshortDecaysInVelo, bool RemoveMultiple, bool weighted,string sExtraVar){

    gROOT->SetBatch(); //ROOT stops plotting canvases

    string magnet = "both";
    bool useRefChannel = !PHSP && !UseOnlyMuMuEvents;
    string MCtype = useRefChannel ? "Reference channel" : (PHSP ? "PHSP" : "signal MC");

    coutInfo("Get efficiencies from " + MCtype + ".");
    TGraphErrors *effGraph = new TGraphErrors();

    coutInfo("Get efficiencies from signal channel in bins of " + sExtraVar + ".");
    int nBins = 0;
    TMefficiencyClass extraVar  = TMefficiencyClass(sExtraVar);
    if (sExtraVar != 0){
        nBins = extraVar.Bins;
    }
    double binCenter = 0;
    double binError = 0;
    if (nBins ==0){
        coutERROR("Wrong variable used!");
        return 0; //I'm feeling adventerous
    }

    vector <double> binBoundaries = extraVar.isEquidistant ? extraVar.binEdgesEquidistant : extraVar.binEdges;
    for (int bin = 0; bin < nBins; bin++){
        EffAndError efficiency = EffAndError(); //put zeroes everywhere
        efficiency = getBDTEfficiencySimple(year,TMVAcut,Run,UseOnlyMuMuEvents,PHSP,KshortDecaysInVelo,RemoveMultiple,weighted,sExtraVar,bin);
        binCenter = (binBoundaries.at(bin+1)+binBoundaries.at(bin))/2.0;
        binError  = (binBoundaries.at(bin+1)-binBoundaries.at(bin))/2.0;
        if (sExtraVar == "q2_binned"){ //Gev
            binCenter = binCenter/1e6;
            binError = binError/1e6;
        }
        effGraph->SetPoint(effGraph->GetN(), binCenter, efficiency.value);
        effGraph->SetPointError(effGraph->GetN()-1,binError,efficiency.highError); //Add error if necessary
        coutDebug("---" + year + " efficiency in bin " + to_string(bin) + ": " + to_string(efficiency.value) + "+-" + to_string(efficiency.highError));
    }
    string name = "BDTeffScan_";
    if (Run ==0) name.append(year);
    else name.append("Run"+to_string(Run));

    //Write to file and close
    string path =  GetEfficiencyFile("BDT", year, magnet, Run, useRefChannel, UseOnlyMuMuEvents, PHSP, KshortDecaysInVelo, RemoveMultiple, weighted, false, false, sExtraVar);
    coutDebug("Writting graph into "+ path);

    TFile * effFile = new TFile(path.c_str(), "RECREATE");
    effFile->cd();
    effGraph->Write("",TObject::kWriteDelete);


    TCanvas * c1 = new TCanvas("c1", "c1");
    c1->cd();

    //Make the plot pretty and plot it
    designBasicEfficiency(effGraph,name,extraVar.sVariable,"{"+string(TMVAmethod)+"}",0.2,0.6); //Add error if necessary
    effGraph->Draw("ap");
    addYearTag(0.65,0.3,year, Run, 1, 1.42);
    replace(path,".root",".eps");
    c1->SaveAs(path.c_str(),"eps");

    effFile->Close();

    coutInfo("All binned BDT efficiency at " + to_string(TMVAcut) +" in " + sExtraVar + " for " + MCtype + " done.");

    return 1;


}

int plotBDTEffciencyInAngles(string year, int Run){
     double finalTMVAcut = (Run == 0 ? getTMVAcut(getRunID(year)) : getTMVAcut(Run));
     plotBDTEfficiency(year, Run, finalTMVAcut, false, true, false, true, true, "q2_binned");
     plotBDTEfficiency(year, Run, finalTMVAcut, false, true, false, true, true, "thetak_equal");
     plotBDTEfficiency(year, Run, finalTMVAcut, false, true, false, true, true, "thetal_equal");
     plotBDTEfficiency(year, Run, finalTMVAcut, false, true, false, true, true, "phi");
     return 1;
}


//----------------------------------------------------------------------------------------------------------
// Selection efficiency
//----------------------------------------------------
TGraphErrors *getSelectionEffTGraph(bool full,bool UseOnlyMuMuEvents, bool PHSP, bool KshortDecaysInVelo,
                             bool RemoveMultiple, bool weighted,
                             string customTMbrach, bool gammaTM){

    bool useRefChannel = !PHSP && !UseOnlyMuMuEvents;
    string MCtype = useRefChannel ? "ReferenceChannel" : (PHSP ? "PHSP" : "SignalMC");

    coutInfo("Get efficiencies from signal channel from " + MCtype+ ".");
    TGraphErrors *effGraph = new TGraphErrors();
    for (auto& yr : yearsMC(useRefChannel, PHSP,12)){
        EffAndError efficiency = EffAndError(); //put zeroes everywhere
        efficiency = getSelectionEfficiencySimple(full, yr,0, UseOnlyMuMuEvents, useRefChannel, PHSP, KshortDecaysInVelo, RemoveMultiple, weighted, "",-1, customTMbrach, gammaTM);
        effGraph->SetPoint(effGraph->GetN(), stoi(yr), efficiency.value);
        effGraph->SetPointError(effGraph->GetN()-1,0.5,efficiency.highError); //Possibly add error
        coutDebug("---"+yr+" efficiency: "+to_string(efficiency.value)+"+-"+to_string(efficiency.highError));
    }
    Color_t color = kBlack;
    if (PHSP) color = kBlue +1;
    if (useRefChannel) color = kRed;
    designBasicEfficiency(effGraph,"SelEffScan"+MCtype,"Year","{Selection}",0.0,full ? 0.00075 : 0.002,color);
    if (!full) effGraph->GetYaxis()->SetTitleOffset(1.75);
    return effGraph;

}

TGraphErrors *getSelectionEffTGraph(bool full, string year, int Run, string customTMbranch, bool gammaTM, bool UseOnlyMuMuEvents, bool Reference, bool PHSP, bool RemoveMultiple, bool weighted, string sExtraVar){

     string MCtype = Reference ? "ReferenceChannel" : (PHSP ? "PHSP" : "SignalMC");
     bool KshortDecaysInVelo = false;
     coutInfo("Get efficiencies from " + MCtype + ".");
     TGraphErrors *effGraph = new TGraphErrors();

     coutInfo("Get efficiencies from signal channel in bins of " + sExtraVar + ".");
     int nBins = 0;
     TMefficiencyClass extraVar  = TMefficiencyClass(sExtraVar);
     if (sExtraVar != 0){
         nBins = extraVar.Bins;
     }
     double binCenter = 0;
     double binError = 0;
     if (nBins ==0){
         coutERROR("Wrong variable used!");
         return 0; //I'm feeling adventerous
     }

     vector <double> binBoundaries = extraVar.isEquidistant ? extraVar.binEdgesEquidistant : extraVar.binEdges;
     for (int bin = 0; bin < nBins; bin++){
         binCenter = (binBoundaries.at(bin+1)+binBoundaries.at(bin));
         binError  = (binBoundaries.at(bin+1)-binBoundaries.at(bin));

         if (sExtraVar == "q2_binned"){
            binCenter = binCenter/2.0e6;
            binError  = binError/2.0e6;
         }

         EffAndError efficiency = EffAndError(); //put zeroes everywhere
         efficiency = getSelectionEfficiencySimple(full, year, Run, UseOnlyMuMuEvents, Reference, PHSP, KshortDecaysInVelo, RemoveMultiple, weighted, sExtraVar, bin, customTMbranch, gammaTM);
         effGraph->SetPoint(effGraph->GetN(), binCenter, efficiency.value);
         effGraph->SetPointError(effGraph->GetN()-1,binError,efficiency.highError); //Add error if necessary
         coutDebug("---" + year + " efficiency in bin " + to_string(bin) + ": " + to_string(efficiency.value) + "+-" + to_string(efficiency.highError));
     }

     Color_t color = kBlack;
     if (PHSP) color = kBlue +1;
     if (Reference) color = kRed;

     string name = "TMeffScan_";
     if (Run ==0) name.append(year);
     else name.append("Run"+to_string(Run));
     designBasicEfficiency(effGraph,name,extraVar.sVariable,"{Selection}",0.0,PHSP ? 0.1 : 0.0075);
     return effGraph;
}

int plotSelectionEfficiency(bool full, string year, int Run,  bool UseOnlyMuMuEvents, bool Reference, bool PHSP, bool weighted, string sExtraVar, string customTMbranch, bool gammaTM){

    gROOT->SetBatch(); //ROOT stops plotting canvases

    bool KshortDecaysInVelo = false;
    string MCtype = Reference ? "Reference channel" : (PHSP ? "PHSP" : "signal MC");
    bool RemoveMultiple = false;

    coutInfo("Get efficiencies from " + MCtype + ".");

    TGraphErrors *effGraph;
    if (sExtraVar == "") effGraph = getSelectionEffTGraph(full,UseOnlyMuMuEvents,PHSP,KshortDecaysInVelo,RemoveMultiple,weighted,customTMbranch,gammaTM);
    else effGraph = getSelectionEffTGraph(full, year, Run, customTMbranch, gammaTM, UseOnlyMuMuEvents, Reference, PHSP, RemoveMultiple, false, sExtraVar);

    //Write to file and close
    string path = GetEfficiencyFileTM("Selection", year, "both", Run, Reference, UseOnlyMuMuEvents, PHSP, KshortDecaysInVelo, RemoveMultiple, weighted, false, "", false, false, sExtraVar, false, customTMbranch, gammaTM);
    if (sExtraVar == "") path = GetEfficiencyFile("Selection","2011",12,Reference,UseOnlyMuMuEvents,PHSP,KshortDecaysInVelo,RemoveMultiple,false,false,false,"");

    coutInfo("Writting graph into " + path);
    if (full) replace(path,".root","_full.root");


    TFile * effFile = new TFile(path.c_str(), "RECREATE");
    effFile->cd();
    effGraph->Write("",TObject::kWriteDelete);


    TCanvas * c1 = new TCanvas("c1", "c1");
    c1->SetTopMargin(0.1);
    c1->SetLeftMargin(0.16);
    c1->cd();
    double yrange = full ? 0.0025 : 0.015;
    if (PHSP) yrange =  full ? 0.0025 : 0.02;
    effGraph->GetYaxis()->SetRangeUser(0.0,yrange);
    effGraph->GetYaxis()->SetNdivisions(505);
    effGraph->GetYaxis()->SetTitleOffset(1.2);
    effGraph->Draw("ap");
    effGraph->GetYaxis()->SetRangeUser(0.0,yrange);
    addYearTag(0.65,0.3,year, Run, 1, 1.42);

    replace(path,".root",".eps");
    c1->SaveAs(path.c_str(),"eps");

    effFile->Close();

    coutInfo("All binned Selection efficiency  in "+sExtraVar+" for "+MCtype+" done.");

    return 1;

}

int plotSelectionEfficiencyAllYearsAll(bool full, bool RemoveMultiple, bool weighted, string customTMbrach, bool gammaTM ){

    gROOT->SetBatch(); //ROOT stops plotting canvases

    bool KshortDecaysInVelo = false;
    coutInfo("Get efficiencies from signal MC, reference channel and PHSP.");

    //Create TGraphErrors
    TGraphErrors *effGraphMuMu = getSelectionEffTGraph(full, true, false, KshortDecaysInVelo, RemoveMultiple, weighted, customTMbrach, gammaTM);
    TGraphErrors *effGraphJpsi = getSelectionEffTGraph(full, false, false, KshortDecaysInVelo, RemoveMultiple, weighted, customTMbrach, gammaTM);
   // TGraphErrors *effGraphPHSP = getSelectionEffTGraph(full, false, true, KshortDecaysInVelo, RemoveMultiple, weighted, customTMbrach, gammaTM);

    //Write to file and close
    string path =  GetEfficiencyFile("Selection", "", "both", 12, false, false, false, KshortDecaysInVelo, RemoveMultiple,weighted, false, false, "",customTMbrach,gammaTM);
    //replace(path,"Run","All_Run");
    if (full) replace(path,".root","_full.root");
    coutInfo("Writting graph into "+ path);

    TFile * effFileTM = new TFile(path.c_str(), "RECREATE");
    effFileTM->cd();
    effGraphMuMu->Write("",TObject::kWriteDelete);
    effGraphJpsi->Write("",TObject::kWriteDelete);
    //effGraphPHSP->Write("",TObject::kWriteDelete);

    TCanvas * c1 = new TCanvas("c1", "c1");
    c1->cd();
    c1->SetLeftMargin(0.16);
    if (full) c1->SetTopMargin(0.1);
    TMultiGraph *mg = new TMultiGraph();
    mg->Add(effGraphMuMu);
    mg->Add(effGraphJpsi);
    //mg->Add(effGraphPHSP);

    mg->Draw("ap");
    mg->GetXaxis()->SetNdivisions(110);
    double yrange = full ? 0.003 : 0.025;
    designMultiGraphEfficiency(mg,"Selection","Year","{Selection}", 0.0, yrange);
    mg->GetYaxis()->SetTitleOffset(1.75);

    mg->Draw("ap");

    TLegend *leg = new TLegend(0.83,full ? 0.87 : 0.9 ,0.6, full ? 0.65 : 0.73);
    leg->AddEntry(effGraphMuMu, "Signal","l");
    leg->AddEntry(effGraphJpsi, "Reference","l");
    leg->Draw("SAME");

    effFileTM->cd();
    c1->Write();
    leg->Write();

    replace(path,".root",".eps");
    c1->SaveAs(path.c_str(),"eps");

    effFileTM->Close();
    return 1;
}

int plotSelectionEfficiencyAll(bool full, bool UseOnlyMuMuEvents, bool Reference, bool PHSP, bool weighted, string sExtraVar){
    for (auto yr: yearsMC(Reference,PHSP,12)){
        if (plotSelectionEfficiency(full, yr, 0, UseOnlyMuMuEvents, Reference, PHSP, weighted, sExtraVar, "TMed", gammaTMdefault) ==0) return 0;
    }
    if (plotSelectionEfficiency(full, "2011", 1, UseOnlyMuMuEvents, Reference, PHSP, weighted, sExtraVar, "TMed", gammaTMdefault)==0) return 0;
    if (plotSelectionEfficiency(full, "2015", 2, UseOnlyMuMuEvents, Reference, PHSP, weighted, sExtraVar, "TMed", gammaTMdefault)==0) return 0;
    return 1;
}
int plotSelectionEfficiencyAll(bool full){
    if (plotSelectionEfficiencyAll(full, false, false, false, true, "q2_binned") ==0) return 0;
    if (plotSelectionEfficiencyAll(full, false, true, false, true, "q2_binned") ==0) return 0;
    if (plotSelectionEfficiencyAll(full, false, false, true, true,"q2_binned") ==0) return 0;
    return 1;
}


//----------------------------------------------------------------------------------------------------------
//Get TM efficiency
//----------------------------------------------------------------------------------------------------------

//Get overall TM efficiency from B+ mass
//----------------------------------------------------
TGraphErrors *getTMeffTGraph(bool UseOnlyMuMuEvents, bool PHSP, bool KshortDecaysInVelo,
                             bool RemoveMultiple, bool weighted,
                             string customTMbrach, bool gammaTM){

    bool useRefChannel = !PHSP && !UseOnlyMuMuEvents;
    string MCtype = useRefChannel ? "ReferenceChannel" : (PHSP ? "PHSP" : "SignalMC");

    coutInfo("Get efficiencies from signal channel from " + MCtype+ ".");
    TGraphErrors *effGraph = new TGraphErrors();
    for (auto& yr : yearsMC(false, false,12)){
        EffAndError efficiency = EffAndError(); //put zeroes everywhere
        efficiency = getTMEfficiencySimple(yr,0, UseOnlyMuMuEvents, PHSP, KshortDecaysInVelo, RemoveMultiple, weighted, "",-1, customTMbrach, gammaTM);
        effGraph->SetPoint(effGraph->GetN(), stoi(yr), efficiency.value);
        effGraph->SetPointError(effGraph->GetN()-1,0.5,efficiency.highError); //Possibly add error
        coutDebug("---"+yr+" efficiency: "+to_string(efficiency.value)+"+-"+to_string(efficiency.highError));
    }
    Color_t color = kBlack;
    if (PHSP) color = kBlue +1;
    if (useRefChannel) color = kRed;
    designBasicEfficiency(effGraph,"TMeffScan"+MCtype,"Year","{TM}",0.25,1.05,color);
    return effGraph;

}

int plotTruthMatchingEfficiencyAllYearsAll(bool RemoveMultiple, bool weighted, string customTMbrach, bool gammaTM ){

    gROOT->SetBatch(); //ROOT stops plotting canvases

    bool KshortDecaysInVelo = false;
    bool UseLowQ2Range = false;

    coutInfo("Get efficiencies from signal MC, reference channel and PHSP.");

    //Create TGraphErrors
    TGraphErrors *effGraphMuMu = getTMeffTGraph(true, false,KshortDecaysInVelo,RemoveMultiple,weighted,customTMbrach, gammaTM);
    TGraphErrors *effGraphJpsi = getTMeffTGraph(false,false,KshortDecaysInVelo,RemoveMultiple,weighted,customTMbrach, gammaTM);
    TGraphErrors *effGraphPHSP = getTMeffTGraph(false,true, KshortDecaysInVelo,RemoveMultiple,weighted,customTMbrach, gammaTM);

    //Write to file and close
    string path =  GetEfficiencyFile("TM", "", "both", 12, false, false, false, KshortDecaysInVelo, RemoveMultiple,weighted, UseLowQ2Range, false, "",customTMbrach,gammaTM);
    replace(path,"Run","All_Run");
    coutInfo("Writting graph into "+ path);

    TFile * effFileTM = new TFile(path.c_str(), "RECREATE");
    effFileTM->cd();
    effGraphMuMu->Write("",TObject::kWriteDelete);
    effGraphJpsi->Write("",TObject::kWriteDelete);
    effGraphPHSP->Write("",TObject::kWriteDelete);

    TCanvas * c1 = new TCanvas("c1", "c1");
    c1->cd();
    TMultiGraph *mg = new TMultiGraph();
    mg->Add(effGraphMuMu);
    mg->Add(effGraphJpsi);
    mg->Add(effGraphPHSP);
    mg->Draw("ap");
    mg->GetXaxis()->SetNdivisions(110);
    designMultiGraphEfficiency(mg,"TMeffScan","Year","{TM}", 0.5, 1.05);
    mg->Draw("ap");
    replace(path,".root",".eps");
    c1->SaveAs(path.c_str(),"eps");

    effFileTM->Close();
    return 1;
}
int plotTruthMatchingEfficiencyAllYearsAll(bool gammaTM, string customTMbranch){
    if (plotTruthMatchingEfficiencyAllYearsAll(false,false, customTMbranch, gammaTM)==0)return 0;
    if (plotTruthMatchingEfficiencyAllYearsAll(false,true,  customTMbranch, gammaTM)==0)return 0;
    if (plotTruthMatchingEfficiencyAllYearsAll(true, false, customTMbranch, gammaTM)==0)return 0;
    if (plotTruthMatchingEfficiencyAllYearsAll(true, true,  customTMbranch, gammaTM)==0)return 0;
    return 1;
}

//Get TM efficiency from B+ mass in bins of sExtraVar
//----------------------------------------------------
TGraphErrors *getTMeffTGraph(string year, int Run, string customTMbranch, bool gammaTM, bool UseOnlyMuMuEvents, bool PHSP, bool RemoveMultiple, bool weighted, string sExtraVar){

     bool useRefChannel = !PHSP && !UseOnlyMuMuEvents;
     string MCtype = useRefChannel ? "ReferenceChannel" : (PHSP ? "PHSP" : "SignalMC");
     bool KshortDecaysInVelo = false;
     coutInfo("Get efficiencies from " + MCtype + ".");
     TGraphErrors *effGraph = new TGraphErrors();

     coutInfo("Get efficiencies from signal channel in bins of " + sExtraVar + ".");
     int nBins = 0;
     TMefficiencyClass extraVar  = TMefficiencyClass(sExtraVar);
     if (sExtraVar != 0){
         nBins = extraVar.Bins;
     }
     double binCenter = 0;
     if (nBins ==0){
         coutERROR("Wrong variable used!");
         return 0; //I'm feeling adventerous
     }

     vector <double> binBoundaries = extraVar.isEquidistant ? extraVar.binEdgesEquidistant : extraVar.binEdges;
     for (int bin = 0; bin < nBins; bin++){
         EffAndError efficiency = EffAndError(); //put zeroes everywhere
         efficiency = getTMEfficiencySimple(year,Run,UseOnlyMuMuEvents,PHSP,KshortDecaysInVelo,RemoveMultiple,weighted,sExtraVar,bin,customTMbranch, gammaTM);
         binCenter = (binBoundaries.at(bin+1)+binBoundaries.at(bin))/2.0;
         effGraph->SetPoint(effGraph->GetN(), binCenter, efficiency.value);
         effGraph->SetPointError(effGraph->GetN()-1,(binBoundaries.at(bin+1)-binBoundaries.at(bin))/2.0,efficiency.highError); //Add error if necessary
         coutDebug("---" + year + " efficiency in bin " + to_string(bin) + ": " + to_string(efficiency.value) + "+-" + to_string(efficiency.highError));
     }

     Color_t color = kBlack;
     if (PHSP) color = kBlue +1;
     if (useRefChannel) color = kRed;

     string name = "TMeffScan_";
     if (Run ==0) name.append(year);
     else name.append("Run"+to_string(Run));
     designBasicEfficiency(effGraph,name,extraVar.sVariable,"{TM}",0.25,1.05);
     return effGraph;
}

int plotTruthMatchingEfficiency(string year, int Run,  bool UseOnlyMuMuEvents, bool PHSP, string sExtraVar, string customTMbranch, bool gammaTM){ //TODO sigEff

    gROOT->SetBatch(); //ROOT stops plotting canvases

    string magnet = "both";
    bool KshortDecaysInVelo = false;
    bool useRefChannel = !PHSP && !UseOnlyMuMuEvents;
    string MCtype = useRefChannel ? "Reference channel" : (PHSP ? "PHSP" : "signal MC");

    coutInfo("Get efficiencies from " + MCtype + ".");
    TGraphErrors *effGraph = getTMeffTGraph(year,Run, customTMbranch , gammaTM, UseOnlyMuMuEvents,PHSP,false,false, sExtraVar); //TODO customTMbranch

    //Write to file and close
    string path = GetEfficiencyFileTM("TM", year, magnet, Run, useRefChannel, UseOnlyMuMuEvents, PHSP, KshortDecaysInVelo, false, false,false, "", false, false, sExtraVar, false, customTMbranch, gammaTM);
    coutInfo("Writting graph into " + path);

    TFile * effFile = new TFile(path.c_str(), "RECREATE");
    effFile->cd();
    effGraph->Write("",TObject::kWriteDelete);

    TCanvas * c1 = new TCanvas("c1", "c1");
    c1->cd();
    effGraph->Draw("ap");
    replace(path,".root",".eps");
    c1->SaveAs(path.c_str(),"eps");

    effFile->Close();

    coutInfo("All binned TM efficiency  in "+sExtraVar+" for "+MCtype+" done.");

    return 1;

}

//----------------------------------------------------------------------------------------------------------
// Run all the stuff for nice plots
//----------------------------------------------------------------------------------------------------------

int runAllEff(){

    //Selection efficiency
    plotSelectionEfficiencyAllYearsAll(false,false,true,"",gammaTMdefault);
    plotSelectionEfficiencyAll(false,true,false,false,true,"q2_binned");
    plotSelectionEfficiencyAll(false,false,true,false,true,"q2_binned");
    //BDT
    ScanBDTEfficiencyAll();
    return 1;
}