//Functions that calculate/fit the yield for given MVA response cut
//and make pretty plots from it
//Renata Kopecna

#include "GlobalFunctions.hh"
#include "Paths.hpp"
#include "MassFit.hpp"
#include "Utils.hpp"
#include "./BmassShape/SignalType.hpp"
#include "EfficiencyClass.cpp"


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

class YieldInfo{
 public:
    string year;
    double sigYield;
    double sigYieldErr;
    double refYield;
    double refYieldErr;
    double bkgYield;
    double bkgYieldErr;
    double sigEff;
    double refEff;
    double allSigEvts;
    YieldInfo(){ //default constructor
        year        = "";
        sigYield    = 0;
        sigYieldErr = 0;
        refYield    = 0;
        refYieldErr = 0;
        bkgYield    = 0;
        bkgYieldErr = 0;
        sigEff      = 0;
        refEff      = 0;
        allSigEvts  = 0;
    }
    void addYield(YieldInfo addInfo){
        sigYield    += addInfo.sigYield;
        sigYieldErr += addInfo.sigYieldErr;
        refYield    += addInfo.refYield;
        refYieldErr += addInfo.refYieldErr;
        bkgYield    += addInfo.bkgYield;
        bkgYieldErr += addInfo.bkgYieldErr;
        sigEff      += addInfo.sigEff; //this doesn't make much sense
        refEff      += addInfo.refEff; //this doesn't make much sense
        allSigEvts  += addInfo.allSigEvts;
    }

    ~YieldInfo();   //destuctor
};

YieldInfo::~YieldInfo(){//destuctor

}

bool FixShape = true;

bool RemoveMultiple = true; //For now use as a global boolean

//Function that takes upper mass sideband histogram and fits it with an exponential
//Using a Tree, starting the upper mass sideband from range+PDGMASS.Bplus and possibly cutting away stuff defined by cut
double GetBackgroundFromSidebandFit(TChain *Tree, double range, string cut, string path){

    TString massBranch = UseDTF ? "B_plus_M" : "B_plus_DTF_M";
    bool savePlots = true;

    //Fill histogram from Tree
    int nBins = 50;
    TH1D *h_B_mass = new TH1D("h_B_mass","h_B_mass",nBins,PDGMASS.B_PLUS + range, cut_B_plus_M_high);
    double binWidth = (cut_B_plus_M_high - PDGMASS.B_PLUS - range)/nBins;
    Tree->Draw(massBranch+" >> h_B_mass",cut.c_str());

    //Define fit function
    TF1 *f_bkg = new TF1("f_bkg","expo",cut_B_plus_M_low,cut_B_plus_M_high);
    design_function(f_bkg,kRed,1);

    //Fit
    h_B_mass->Fit(f_bkg,"Q","", PDGMASS.B_PLUS + range,cut_B_plus_M_high); //exponential is f(x) = exp(p0+p1*x)
    double integral = f_bkg->Integral(PDGMASS.B_PLUS-range,PDGMASS.B_PLUS+range)/binWidth;

    //Save the plots
    if (savePlots){
        string name = path.substr(path.find("/Kplus")+6,path.length());
        //Create canvas
        TCanvas  *c = c_canvas(name);
        c->cd();
        design_TH1D(h_B_mass,"B^{+} mass [MeV]","Counts a.u.", kBlack);
        h_B_mass->GetYaxis()->SetTitleOffset(1.05);
        h_B_mass->Draw();
        //Create TPaveText
        TPaveText *fitParams= new TPaveText(0.6,0.7,0.87,0.9,"NDC");
        fitParams->SetFillColor(kWhite);
        fitParams->AddText(Form("Offset : %.2f ",  f_bkg->GetParameter(0)));
        fitParams->AddText(Form("Exp: %.2f%%", f_bkg->GetParameter(1)*100));
        fitParams->AddText(Form("Integral: %.2f ", integral));
        fitParams->Draw("SAME");

        //Create file
        TFile *fitFile  = new TFile(path.c_str(),"RECREATE");
        coutDebug("Opening " + string(fitFile->GetPath()));
        fitFile->cd();
        h_B_mass->Write();
        f_bkg->Write();
        c->Write();
        replace(path,".root",".eps");
        c->SaveAs(path.c_str(), ".eps");
        fitFile->Close();
    }
    //return the integral
    h_B_mass->Clear();
    return integral;

}

double GetMVAefficiency(string year, int Run, bool KshortDecaysInVelo,
                        bool UseLowQ2Range, bool IncludeMultipleEff, bool sig, double TMVAcut){
    //Load the efficiencies
    coutDebug("Openning MVA efficiency files.");
    string path = GetEfficiencyFile("BDT", year, "both", Run, !sig, sig,  false, KshortDecaysInVelo, IncludeMultipleEff, true, UseLowQ2Range, false,"");

    if (Run == 0 && std::stoi(year) > 2016 && !sig) replace(path,year,"2016"); //RefChann available for 2015-2016
    if (Run == 0 && year=="2015" && sig) replace(path,year,"2016");            //No reasonable MC for 2015

    coutDebug("Openning file: " + path);
    TFile * effFile = new TFile(path.c_str(), "OPEN");

    //Get the TGraphs
    coutDebug("File opened, getting the graphs now.");
    string graphName = "effScan_" + (Run ==0 ? year : "Run" + to_string(Run));
    if (Run == 0 && std::stoi(year) > 2016) replace(graphName,year,"2016");
    if (Run == 0 && year=="2015" && sig)    replace(graphName,year,"2016");
    coutDebug("\t" + graphName);

    TGraphErrors *effGraph= (TGraphErrors*)effFile->Get(TMVAmethod+graphName.c_str());
    double efficiency = effGraph->Eval(TMVAcut);

    coutDebug(string(sig ? "Signal   " : "Reference") +" MVA efficiency: " + to_string(efficiency));

    //Close the efficiency files
    effFile->Close();
    return efficiency;

}

double GetSelectionEfficiency(string year, int Run, bool KshortDecaysInVelo,
                              bool UseLowQ2Range, bool RemoveMultiple, bool sig){
    //Load the efficiencies
    coutDebug("Openning efficiency files.");
    EffAndError eff = getSelectionEfficiencySimple(true,year,Run,sig,!sig, false, KshortDecaysInVelo, RemoveMultiple, true, "", -1, "", gammaTMdefault);

    coutDebug(string(sig ? "Signal   " : "Reference") +" selection efficiency: " + to_string(eff.value));

    //Close the efficiency files
    return eff.value;
}


string GetBackgroundFunction(bool KshortDecaysInVelo){    //Set signal and background function separately for each decay channel
    string BackGroundFunction = "SingleExponential";
    if(Kst2Kspiplus){
        if(KshortDecaysInVelo) BackGroundFunction = "SingleExponential";
        else                   BackGroundFunction = "DoubleExponential";
    }
    return BackGroundFunction;
}

//Takes yield from the reference channel,
YieldInfo GetSigAndBkgEstimation(string year, int Run, bool KshortDecaysInVelo,
                                 bool UseLowQ2Range, Double_t TMVAcut, bool scan){

    bool fixedMassRegion = !Kst2Kspiplus; //Calculates the signal from +- 100 MeV or +- SignalRegionNsigma*effective sigma

    //initialize the struct
    YieldInfo yield = YieldInfo();
    yield.year = year;

    bool RemoveMultiple = true;

    //Set signal and background function separately for each decay channel
    string SignalFunction = "OneCB";
    string BackGroundFunction = GetBackgroundFunction(KshortDecaysInVelo);

    //Get the reference signal yield
    if (!scan){
        coutInfo("Running the fit!");
        yield.refYield = basicYieldFit(year,Run,
                                       false, false, true, false,
                                       FixShape, SignalFunction, BackGroundFunction, FixShape,
                                       KshortDecaysInVelo, UseLowQ2Range,
                                       TMVAcut, fixedMassRegion, true, RemoveMultiple);
    }

    //Load the file from fitting
    TFile *fitFile  = new TFile(GetMassFitFile(year, Run,
                                               false, false, true, false,
                                               FixShape, SignalFunction, BackGroundFunction, FixShape,
                                               KshortDecaysInVelo, UseLowQ2Range,
                                               TMVAcut, fixedMassRegion, RemoveMultiple).c_str(),"OPEN");

    coutDebug("Opening " + string(fitFile->GetPath()));

    if (scan) yield.refYield = getSigYield(fitFile);
    double effSigma = fixedMassRegion ? B_plus_M_signal_window : getEffSigma(fitFile);

    if (yield.refYield == 0) coutWarning("Reference channel yield is zero!");

    yield.sigEff = GetMVAefficiency(year,Run,KshortDecaysInVelo,UseLowQ2Range,RemoveMultiple,true,TMVAcut) * GetSelectionEfficiency(year, Run, KshortDecaysInVelo, UseLowQ2Range, !RemoveMultiple, true);
    yield.refEff = GetMVAefficiency(year,Run,KshortDecaysInVelo,UseLowQ2Range,RemoveMultiple,false,TMVAcut) * GetSelectionEfficiency(year, Run, KshortDecaysInVelo, UseLowQ2Range, !RemoveMultiple, false);

    coutDebug("Sig Yield efficiency = " + to_string(yield.sigEff));
    coutDebug("Ref Yield efficiency = " + to_string(yield.refEff));

    //Calculate the signal yield from reference yield, efficiencies and the Jpsi->mumu BR
    yield.sigYield = BR_sig / BR_ref * yield.refYield * yield.sigEff/yield.refEff;

    //Get the background yield from side band
    TChain * tree   = new TChain("DecayTree");

    std::vector<string> years;
    if (Run == 0) years.push_back(year);
    else years = yearsData(Run);

    for (auto& yr : years){
        string BDToutputPath = GetBDToutputFile(yr,getRunID(yr),false,false,false,KshortDecaysInVelo,UseLowQ2Range,false);
        tree->Add(BDToutputPath.c_str());
        coutDebug("Opening " + BDToutputPath);
    }

    string sVariable = (UseDTF ? "B_plus_M_DTF" : "B_plus_M");
    string cutMass = Form("%s > %f && %s < %f",
                           sVariable.c_str(), getBplusMeanFromResult(fitFile) -  effSigma, sVariable.c_str(),
                           getBplusMeanFromResult(fitFile) + effSigma);
    coutDebug("B mass cut: " + cutMass);

    //cut away Jpsi
    string cut = getFinalCut(false, false, false, "", gammaTMdefault, true, false, false, false, false, TMefficiencyClass(), -1, TMVAcut, RemoveMultiple);
    coutDebug("Using cut " + cut);

    //Get background estimation from sideband
    string bkgFitPath = GetBDTScanBackgroundFitFile(year,Run,KshortDecaysInVelo,UseLowQ2Range,TMVAcut);
    yield.bkgYield = GetBackgroundFromSidebandFit(tree, fixedMassRegion ? effSigma : B_plus_M_signal_window, cut, bkgFitPath);
    cut = cut + " && (" + cutMass + ")";
    coutDebug("Using cut " + cut);

    //data: create histograms from tree
    tree->Draw(Form("%s>>B_plus_M_plot", sVariable.c_str()), cut.c_str());
    TH1 * histo = ((TH1 *) gPad->GetPrimitive("B_plus_M_plot"));
    Int_t EventsAfterPreSelection = histo->GetEntries();

    yield.allSigEvts = EventsAfterPreSelection;

    //Print signal, background and all evts
    string trackType = (Kst2Kspiplus && SplitDDandLL ? (KshortDecaysInVelo ? "LL Tracks: " : "DD Tracks: ") : ":");
    coutDebug(TheDecay + "[SIGNAL]: "  + trackType + to_string(yield.sigYield));
    coutDebug(TheDecay + "[BCKGND]: "  + trackType + to_string(yield.bkgYield));
    coutDebug(TheDecay + "[ALLEVTS]: " + trackType + to_string(yield.allSigEvts));

    coutInfo("Expected yield is: " + to_string(yield.sigYield) + "\t+-" + to_string(yield.sigYieldErr));

    return yield;
}

YieldInfo GetSigAndBkgEstimationFromData(string year = "2011", int Run = 0, int randomSubset = 0,
                                         bool KshortDecaysInVelo = false, bool UseLowQ2Range = false,
                                         Double_t TMVAcut = -1.0, bool scan = false){
    gStyle->SetOptStat(0);
    LHCbStyle();

    bool fixedMassRegion = !Kst2Kspiplus; //Fix it in pi0 channel //TODO
    string magnet = "both"; //prepare in case of polarity studies

    //initialize the struct
    YieldInfo yield = YieldInfo();
    yield.year = year;

    //Set signal and background function separately for each decay channel
    string SignalFunction = "OneCB";
    string BackGroundFunction = GetBackgroundFunction(KshortDecaysInVelo);

    if (!scan) { //Do the fit first
        yield.refYield = massFit(year,magnet,Run,
                                 false, true, false, false,
                                 true, false,
                                 FixShape, SignalFunction, BackGroundFunction, FixShape,
                                 KshortDecaysInVelo, UseLowQ2Range,
                                 TMVAcut, randomSubset,
                                 fixedMassRegion, false,
                                 false, true, false,
                                 "", -1,
                                 RemoveMultiple, true,
                                 false, "", false,
                                 "", gammaTMdefault, false);
        yield.sigYield = massFit(year,magnet,Run,
                                 false, true, false, false,
                                 false, true,
                                 FixShape, SignalFunction, BackGroundFunction, FixShape,
                                 KshortDecaysInVelo, UseLowQ2Range,
                                 TMVAcut, randomSubset,
                                 fixedMassRegion, false,
                                 false, true, false,
                                 "", -1,
                                 RemoveMultiple, true,
                                 false, "", false,
                                 "", gammaTMdefault, false);
    }    
    //Load fit results now:

    // Get reference yield
    TFile *fitFileRef  = new TFile(GetMassFitFile(year, magnet, Run,
                                                  false, true, false, false,
                                                  true, false,
                                                  FixShape, SignalFunction, BackGroundFunction, FixShape,
                                                  KshortDecaysInVelo, UseLowQ2Range,
                                                  TMVAcut, randomSubset,
                                                  fixedMassRegion, false,
                                                  "", -1,
                                                  RemoveMultiple, false, //Data cannot be weighted technically, so the name is without "weighted"
                                                  false, "", false,
                                                  "", gammaTMdefault, false).c_str(),"OPEN");

    coutDebug("Opening " + string(fitFileRef->GetPath()));
    yield.refYield    = getSigYield(fitFileRef);
    yield.refYieldErr = getSigYieldErr(fitFileRef);
    fitFileRef->Close();

    // Get signal yield
    TFile *fitFile  = new TFile(GetMassFitFile(year, magnet, Run,
                                               false, true, false, false,
                                               false, true,
                                               FixShape, SignalFunction, BackGroundFunction, FixShape,
                                               KshortDecaysInVelo, UseLowQ2Range,
                                               TMVAcut, randomSubset,
                                               fixedMassRegion, false,
                                               "", -1,
                                               RemoveMultiple, false, //Data cannot be weighted technically, so the name is without "weighted"
                                               false, "", false,
                                               "", gammaTMdefault, false).c_str(),"OPEN");

    coutDebug("Opening " + string(fitFile->GetPath()));
    yield.sigYield    = getSigYield(fitFile);
    yield.sigYieldErr = getSigYieldErr(fitFile);
    yield.bkgYield    = getBkgYield(fitFile);
    yield.bkgYieldErr = getBkgYieldErr(fitFile);
    double mean_sig = getBplusMeanFromResult(fitFile);
    double effSigma = fixedMassRegion ? B_plus_M_signal_window : getEffSigma(fitFile);
    fitFile->Close();

    //Get the number of all events in the signal region
    TChain * tree   = new TChain("DecayTree");
    std::vector<string> years;
    if (Run == 0) years.push_back(year);
    else years = yearsData(Run);

    for (auto& yr : years){
        string BDToutputPath = GetBDToutputFile(yr,getRunID(yr),false,false,false,KshortDecaysInVelo,UseLowQ2Range,false);
        tree->Add(BDToutputPath.c_str());
        coutDebug("Opening " + BDToutputPath);
    }

    coutDebug("eff sigma: " + to_string(effSigma));
    string sVariable = (UseDTF ? "B_plus_M_DTF" : "B_plus_M");
    string cutMass = Form("%s > %f && %s < %f",
                           sVariable.c_str(), mean_sig -  effSigma, sVariable.c_str(),
                          mean_sig + effSigma);
    coutDebug("B mass cut: " + cutMass);

    //cut away resonances
    string cut = getFinalCut(false, false, false, "", gammaTMdefault, true, false, false, false, false, TMefficiencyClass(), -1, TMVAcut, RemoveMultiple);
    coutDebug("Using cut " + cut);

    //data: create histograms from tree
    tree->Draw(Form("%s>>B_plus_M_plot", sVariable.c_str()), cut.c_str());
    TH1 * histo = ((TH1 *) gPad->GetPrimitive("B_plus_M_plot"));
    Int_t EventsAfterPreSelection = histo->GetEntries();

    yield.allSigEvts = EventsAfterPreSelection;

    //Print signal, background and all evts
    string trackType = (Kst2Kspiplus && SplitDDandLL ? (KshortDecaysInVelo ? "LL Tracks: " : "DD Tracks: ") : ":");
    coutDebug(TheDecay + "[SIGNAL]: "  + trackType + to_string(yield.sigYield));
    coutDebug(TheDecay + "[BCKGND]: "  + trackType + to_string(yield.bkgYield));
    coutDebug(TheDecay + "[ALLEVTS]: " + trackType + to_string(yield.allSigEvts));

    coutInfo("Expected yield is: " + to_string(yield.sigYield) + "\t+-" + to_string(yield.sigYieldErr));
    coutDebug("Signal yield error: " + to_string(yield.sigYieldErr));
    coutDebug("Background yield error: " + to_string(yield.bkgYieldErr));
    coutDebug("Reference channel  yield error: " + to_string(yield.refYieldErr));

    if (yield.refYield == 0) coutERROR("Reference channel yield is zero!");
    if (yield.sigYield == 0) coutERROR("Signal channel yield is zero! Abort.");

    return yield;

}

int SaveTGraphs(string path, bool fineScan, string year, int Run, string basicPath,
                TGraphErrors *yieldGraph, TGraphErrors *bkgGraph, TGraphErrors *bkgGraphFromAllEvts, TGraphErrors *refYieldGraph, TGraphErrors *allEvtsInSig, TGraphErrors *significance, TGraphErrors *significanceFromAllEvts){

    coutInfo("Making pretty plots and saving them!");
    coutDebug("Writting into " + path);
    coutDebug("Using basicPath " + basicPath);

    TFile * TGraphOutput = new TFile(path.c_str(), "RECREATE");
    TGraphOutput->cd();

    designYieldGraph(yieldGraph,              Run, year, "sigYield",                 basicPath, TGraphOutput, false);
    designYieldGraph(bkgGraph,                Run, year, "bkgYield",                 basicPath, TGraphOutput, false);
    designYieldGraph(bkgGraphFromAllEvts,     Run, year, "bkgYield_fromAllEvts",     basicPath, TGraphOutput, fineScan);
    designYieldGraph(refYieldGraph,           Run, year, "refYield",                 basicPath, TGraphOutput, fineScan);
    designYieldGraph(significance,            Run, year, "significance",             basicPath, TGraphOutput, false);
    designYieldGraph(significanceFromAllEvts, Run, year, "significance_fromAllEvts", basicPath, TGraphOutput, fineScan);
    designYieldGraph(allEvtsInSig,            Run, year, "EventsInSigRegion",        basicPath, TGraphOutput, false);

    TGraphOutput->Close();
    return 1;


}

int ScanSignalAndBckgndEstimation(string year, int Run, Double_t BDTstep, bool KshortDecaysInVelo, bool UseLowQ2Range, bool scan, bool fineScan){
    //If Run==0, do the scan for year, if Run == 1/2, do the scan per run

    if (verboseLevel > 1)RooMsgService::instance().setGlobalKillBelow(RooFit::FATAL);

    bool section = true; //In case of fine scan, section the scan in more fine way
    int n_divisions = 3;

    //Get low TMVA response lower boundary
    Double_t lowBDTcut = (TMVAmethod == "MLP") ? 0.0 : -1.0;
    if (fineScan){
        coutInfo("Making fine BDT cut study.");
        lowBDTcut = 0.9;
        BDTstep =   0.001;
    }

    //prepare in case of polarity studies
    string magnet = "both";

    TGraphErrors *yieldGraph = new TGraphErrors();
    TGraphErrors *bkgGraph = new TGraphErrors();
    TGraphErrors *bkgGraphFromAllEvts = new TGraphErrors();
    TGraphErrors *refYieldGraph = new TGraphErrors();
    TGraphErrors *significance = new TGraphErrors();
    TGraphErrors *significanceFromAllEvts = new TGraphErrors();
    YieldInfo yield = YieldInfo(); //put zeroes everywhere

    //loop over BDT cuts
    coutInfo("Loop over BDT cuts with a step of " + to_string(BDTstep));
    coutDebug("Starting the loop at " + to_string(lowBDTcut));
    double tmpBDTcut = 0.0;

    for(double fBDTcut = lowBDTcut; fBDTcut < (section ? 0.999 : 1.0); fBDTcut += section ? ((1.0-fBDTcut)/double(n_divisions)) : BDTstep){
        for (int k = 0; k < n_divisions; k++){
            tmpBDTcut = fBDTcut + k*(1.0-fBDTcut)/double(n_divisions);
            yield = GetSigAndBkgEstimation(year, Run, KshortDecaysInVelo, UseLowQ2Range, tmpBDTcut, scan);
            yieldGraph              ->SetPoint(yieldGraph->GetN(),          tmpBDTcut, yield.sigYield);
            bkgGraphFromAllEvts     ->SetPoint(bkgGraphFromAllEvts->GetN(), tmpBDTcut, yield.allSigEvts-yield.sigYield);
            bkgGraph                ->SetPoint(bkgGraph->GetN(),            tmpBDTcut, yield.bkgYield);
            refYieldGraph           ->SetPoint(yieldGraph->GetN(),          tmpBDTcut, yield.refYield);
            significance            ->SetPoint(significance->GetN(),        tmpBDTcut, yield.sigYield / (TMath::Sqrt(yield.sigYield + yield.bkgYield)));
            significanceFromAllEvts ->SetPoint(significance->GetN(),        tmpBDTcut, yield.sigYield / (TMath::Sqrt(yield.allSigEvts)));
            coutDebug("---" + string(TMVAmethod) + " cut " + to_string(tmpBDTcut) + " yield: " + to_string(yield.sigYield) + ", background: " + to_string( yield.allSigEvts-yield.sigYield ));
        }
    }

    //Write to file and close
    string path = GetBDTScanFile(year, magnet, Run, KshortDecaysInVelo, UseLowQ2Range, fineScan);
    SaveTGraphs(path,fineScan,year,Run,path,yieldGraph,bkgGraph,bkgGraphFromAllEvts,refYieldGraph,NULL,significance,significanceFromAllEvts);

    coutInfo("Signal scan done for year " + year + "!");
    return 1;

}

int ScanSignalAndBckgndEstimationPerYear(string year = "2011", Double_t BDTstep = 0.01, bool KshortDecaysInVelo = true, bool UseLowQ2Range = false, bool scan = false, bool fineScan = false){
    return ScanSignalAndBckgndEstimation(year = "2011", 0, BDTstep, KshortDecaysInVelo, UseLowQ2Range, scan, fineScan);
}
int ScanSignalAndBckgndEstimationAllYears(Double_t BDTstep = 0.01, bool KshortDecaysInVelo = true, bool UseLowQ2Range = false, bool scan = false, bool fineScan = true){
    for (auto yr: yearsData(12)) ScanSignalAndBckgndEstimationPerYear(yr, BDTstep, KshortDecaysInVelo, UseLowQ2Range, scan, fineScan);
    return 1;
}

int ScanSignalAndBckgndEstimationSimple(string year = "2011",int Run = 1, Double_t BDTstep = 0.01, int randomSubset = 0,
                                        bool KshortDecaysInVelo = true, bool UseLowQ2Range = false, bool scan = false, bool fineScan = true){
    //Get low TMVA response lower boundary

    double lowBDTcut = (TMVAmethod == "MLP") ? 0.0 : -1.0;
    if (fineScan){
        coutInfo("Making fine BDT cut study.");
        lowBDTcut = 0.9;
        BDTstep = 0.002;
    }
    bool halve = true;


    //prepare in case of polarity studies
    string magnet = "both";

    TGraphErrors *yieldGraph = new TGraphErrors();
    TGraphErrors *bkgGraph = new TGraphErrors();
    TGraphErrors *bkgGraphFromAllEvts = new TGraphErrors();
    TGraphErrors *refYieldGraph = new TGraphErrors();
    TGraphErrors *allEvtsInSig = new TGraphErrors();
    TGraphErrors *significance = new TGraphErrors();
    TGraphErrors *significanceFromAllEvts = new TGraphErrors();
    YieldInfo yield = YieldInfo(); //put zeroes everywhere

    //loop over BDT cuts
    coutInfo("Loop over BDT cuts with a step of " + to_string(BDTstep));
    coutDebug("Starting the loop at " + to_string(lowBDTcut));
    int n_divisions = halve ? 3 : 1;
    double tmpcut = 0.0;
    for(double fBDTcut = lowBDTcut; fBDTcut < (halve ? 0.9999 : 1); fBDTcut += halve ? ((1.0-fBDTcut)/double(n_divisions)) : BDTstep){
        for (int k = 0; k < n_divisions-1; k++){
            tmpcut = fBDTcut + k*(1.0-fBDTcut)/double(n_divisions);
            coutDebug("Cutting at: " + to_string(tmpcut));
            yield =GetSigAndBkgEstimationFromData(year,Run, randomSubset, KshortDecaysInVelo, UseLowQ2Range, tmpcut, scan);
            //Save the efficiency
            yieldGraph              ->SetPoint(yieldGraph->GetN(),          tmpcut, yield.sigYield);
            yieldGraph              ->SetPointError(yieldGraph->GetN()-1,   0,       yield.sigYieldErr);
            bkgGraph                ->SetPoint(bkgGraph->GetN(),            tmpcut, yield.bkgYield);
            bkgGraph                ->SetPointError(bkgGraph->GetN()-1,     0,       yield.bkgYieldErr);
            bkgGraphFromAllEvts     ->SetPoint(bkgGraph->GetN(),            tmpcut, yield.allSigEvts-yield.sigYield);
            refYieldGraph           ->SetPoint(yieldGraph->GetN(),          tmpcut, yield.refYield);
            refYieldGraph           ->SetPointError(yieldGraph->GetN()-1,   0,       yield.refYieldErr);
            allEvtsInSig            ->SetPoint(allEvtsInSig->GetN(),        tmpcut, yield.allSigEvts);
            significance            ->SetPoint(significance->GetN(),        tmpcut, yield.sigYield / (TMath::Sqrt(yield.sigYield+yield.bkgYield)));
            significanceFromAllEvts ->SetPoint(significance->GetN(),        tmpcut, yield.sigYield / (TMath::Sqrt(yield.allSigEvts)));
            coutDebug("---" + string(TMVAmethod) + " cut " + to_string(tmpcut) + " yield: " + to_string(yield.sigYield) + ", background: " + to_string( yield.allSigEvts-yield.sigYield ));
        }
    }

    //Write to file and close //TODO at some point
    string path = GetBDTScanFile(year, magnet, Run, KshortDecaysInVelo, UseLowQ2Range, fineScan);
    replace(path, ".root", "fromData.root");
    if (randomSubset == -1) replace(path,".root","_subset1.root");
    if (randomSubset ==  1) replace(path,".root","_subset2.root");
    SaveTGraphs(path,fineScan,year,0,path,yieldGraph,bkgGraph,bkgGraphFromAllEvts,refYieldGraph,NULL,significance,significanceFromAllEvts);

    coutInfo("Signal scan done for Run " + to_string(Run) + "!");
    return 1;
}

int ScanSignalAndBckgndEstimationSimplePerYear(string year = "2011", Double_t BDTstep = 0.01, int randomSubset = 0, bool KshortDecaysInVelo = true, bool UseLowQ2Range = false, bool scan = false, bool fineScan = true){
    return ScanSignalAndBckgndEstimationSimple(year = "2011", 0, BDTstep, randomSubset, KshortDecaysInVelo, UseLowQ2Range, scan, fineScan);
}

int ScanSignalAndBckgndEstimationSimpleAllYears(Double_t BDTstep = 0.01, int randomSubset = 0, bool KshortDecaysInVelo = true, bool UseLowQ2Range = false, bool scan = false, bool fineScan = true){
    for (auto yr: yearsData(12)) ScanSignalAndBckgndEstimationSimplePerYear(yr, BDTstep, randomSubset, KshortDecaysInVelo, UseLowQ2Range, scan, fineScan);
    return 1;
}

double getMaxBDTresponse(string year = "2011", int Run = 0,  bool fineScan = false, bool directScan = false, int randomSubset = 0, bool KshortDecaysInVelo = true, bool UseLowQ2Range  = false){

    //prepare in case of polarity studies
    string magnet = "both";

    if (Run == 0) coutInfo("Getting BDTresponse with highest significance for YEAR " + year + ". .  .");
    else          coutInfo("Getting BDTresponse with highest significance for RUN "+to_string(Run)+". .  .");

    string path = GetBDTScanFile(year, magnet, Run, KshortDecaysInVelo, UseLowQ2Range, fineScan);
    if (directScan) replace(path, ".root", "fromData.root");
    if (randomSubset == -1) replace(path,".root","_subset1.root");
    if (randomSubset ==  1) replace(path,".root","_subset2.root");

    TGraphErrors *significance = new TGraphErrors();
    TFile *scanFile = new TFile(path.c_str(),"OPEN");
    string graphName = "significance_fromAllEvts";
    significance = (TGraphErrors*) scanFile->Get(graphName.c_str());
    if(significance == NULL){
        coutERROR("TGraphError was not found in the file!");
        coutERROR("File:\t" + path);
        coutERROR("TGraph:\t" + graphName);
        return 0;
    }

    //Get the best BDTcut
    double LargestFoM = -100.;
    double bestCut = 0.0;
    for(int p = 0; p < significance->GetN(); p++){
        if(significance->GetY()[p] > LargestFoM){
            LargestFoM  = significance->GetY()[p];
            bestCut     = significance->GetX()[p];
        }
    }
    return bestCut;
}

int optimizeBDTCut(string year = "2011", int Run = 0, bool KshortDecaysInVelo = true, Double_t BDTstep = 0.01, bool UseRandomSubset = false,
                       bool UseLowQ2Range  = false, bool scan = false, bool fineScan = false, bool directScan = false){

    //Get low TMVA response lower boundary
    double lowBDTcut = (TMVAmethod == "MLP") ? 0.0 : -1.0;
    if (fineScan){
        coutInfo("Making fine BDT cut study.");
        lowBDTcut = 0.9;
        BDTstep = 0.002;
    }

    double BestCut = 0;

    if (directScan){
        if (UseRandomSubset){
            ScanSignalAndBckgndEstimationSimple(year, Run, BDTstep, -1, KshortDecaysInVelo, UseLowQ2Range, scan, fineScan);
            ScanSignalAndBckgndEstimationSimple(year, Run, BDTstep, +1, KshortDecaysInVelo, UseLowQ2Range, scan, fineScan);
            double BestCut1 = getMaxBDTresponse(year, Run, fineScan, true, -1, KshortDecaysInVelo, UseLowQ2Range);
            double BestCut2 = getMaxBDTresponse(year, Run, fineScan, true, +1, KshortDecaysInVelo, UseLowQ2Range);
            BestCut = roundf(BestCut1 * 100) / 100 +roundf(BestCut2 * 100); // round to 2 decimal points, return e.g. 0.96 + 94 (easy to untangle)
        }
        else{
            ScanSignalAndBckgndEstimationSimple(year, Run, BDTstep, 0, KshortDecaysInVelo, UseLowQ2Range, scan, fineScan);
            BestCut = getMaxBDTresponse(year, Run, fineScan, true, 0, KshortDecaysInVelo,  UseLowQ2Range);
        }

        //TODO
    }
    else{ //TODO: add subset fits
        if (Run == 0){
            ScanSignalAndBckgndEstimationPerYear(year,BDTstep,KshortDecaysInVelo,UseLowQ2Range,scan,fineScan);
        }
        else{
            ScanSignalAndBckgndEstimation(year,Run,BDTstep,KshortDecaysInVelo,UseLowQ2Range,scan,fineScan);
        }
        //TODO make a pretty plot

        BestCut = getMaxBDTresponse(year, Run, fineScan, false, 0, KshortDecaysInVelo, UseLowQ2Range);

    }

    return BestCut;

}