//Makes pretty yield (after the MVA cut) plots and tables
//Renata Kopecna

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

class YieldInfo{
 public:
    double sigYield;
    double sigYieldErr;
    double bkgYield;
    double bkgYieldErr;
    double significance;
    YieldInfo(){ //default constructor
        sigYield    = 0;
        sigYieldErr = 0;
        bkgYield    = 0;
        bkgYieldErr = 0;
        significance  = 0;
    }
    void addYield(YieldInfo addInfo){
        sigYield    += addInfo.sigYield;
        sigYieldErr += addInfo.sigYieldErr;
        bkgYield    += addInfo.bkgYield;
        bkgYieldErr += addInfo.bkgYieldErr;
        significance = sigYield ==0 ? 0 : sigYield/sqrt(sigYield+bkgYield);
    }

    ~YieldInfo();   //destuctor
};

YieldInfo::~YieldInfo(){//destuctor

}

double getValueFromTGraph(string name, TFile *fitFile, double TMVAcut){
    //Read the expected yields and significances
    TGraphErrors *graph = (TGraphErrors*)fitFile->Get(name.c_str());
    return graph->Eval(TMVAcut);
}

bool yieldComparison(int Run, double TMVAcut){
    //Takes the expected signal+background yield and significance as well as the actuall yield and significance a couts it per Run
    bool fixedMassRegion = true;


    //Read the expected yields and significances
    TFile *yield = new TFile(GetBDTScanFile("2011","both",Run,false,false,true).c_str(),"READ");
    double expectSigYield = getValueFromTGraph("sigYield",yield,TMVAcut);
    double expectBkgYield = getValueFromTGraph("bkgYield_fromAllEvts",yield,TMVAcut);
    double expectSignificance = getValueFromTGraph("significance_fromAllEvts",yield,TMVAcut);

    basicYieldFit("2011",Run,
                  false,false,false,true,
                  true,"OneCB","SingleExponential", true,
                  false, false,
                  TMVAcut, fixedMassRegion, false, true);
    //Load the file from fitting
    TFile *fitFile  = new TFile(GetMassFitFile("2011", Run,
                                               false, false, false, true,
                                               true, "OneCB", "SingleExponential", true,
                                               false, false,
                                               TMVAcut, fixedMassRegion, true).c_str(),"OPEN");

    double sigYield    = getSigYield(fitFile);
    double sigYieldErr = getSigYieldErr(fitFile);
    double bkgYield    = getBkgYield(fitFile);
    double bkgYieldErr = getBkgYieldErr(fitFile);
    double significance = sigYield / sqrt(sigYield+bkgYield);
    cout << "\\begin{table}[hbt!]" << endl;
    cout << "\\centering" << endl;
    cout << "\\begin{tabular}{l|l|l}" << endl;
    cout << "Run " << Run << " & Expected & Fitted \\\\ \\hline" << endl;
    cout << "Signal &" << expectSigYield << " & " << sigYield << "$\\pm$" << sigYieldErr << "\\\\" << endl;
    cout << "Background &" << expectBkgYield << " & " << bkgYield << "$\\pm$" << bkgYieldErr << "\\\\" << endl;
    cout << "S/sqrt(S+B) &"<< expectSignificance << " & " << significance << endl;
    cout << "\\end{tabular}" << endl;
    cout << "\\end{table}" << endl;

    return 1;
}

YieldInfo yieldInQ2(int Run, double TMVAcut, int nBin, bool doFit){
    //Fits the B mass in diffrent Q2 bins and plots the yields and background

    bool fixedMassRegion = true;
    if (doFit) massFit("2011", "both", Run, false,
            true,true,false,false,true,
            true,"OneCB","SingleExponential", true,
            false, false,
            TMVAcut, false,
            fixedMassRegion, false,
            false, false, false,
            "q2_binned_fit", nBin,
            true,
            false, false, "",
            false, "", gammaTMdefault,  false);

    //Load the file from fitting
     TFile *fitFile  = new TFile(GetMassFitFile("2011", "both", Run, false,
                                                true,true,false,false,true,
                                                true,"OneCB","SingleExponential", true,
                                                false, false,
                                                TMVAcut, false,
                                                fixedMassRegion, false,
                                                "q2_binned_fit", nBin,
                                                true,
                                                false, false, "",
                                                false, "", gammaTMdefault,  false).c_str(),"OPEN");

     YieldInfo YI;
     //Faster to read erything than check million ifs
     YI.sigYield = getSigYield(fitFile);
     YI.sigYieldErr = getSigYieldErr(fitFile);
     YI.bkgYield    = getBkgYield(fitFile);
     YI.bkgYieldErr = getBkgYieldErr(fitFile);
     if (YI.sigYield == 0) YI.significance =0;
     else  YI.significance =  YI.sigYield / sqrt( YI.sigYield+ YI.bkgYield);

    return YI;
}

bool plotYieldInQ2(bool fixRange){


    TMefficiencyClass extraVar  = TMefficiencyClass("q2_binned_fit");

    int nBins = extraVar.Bins;
    double binCenter = 0;
    double binError = 0;
    if (nBins ==0){
        coutERROR("Wrong variable used!");
        return 0; //I'm feeling adventerous
    }

    TGraphErrors *graphSig_1 = new TGraphErrors();
    TGraphErrors *graphBkg_1 = new TGraphErrors();
    TGraphErrors *graphSignificance_1 = new TGraphErrors();
    TGraphErrors *graphSig_2 = new TGraphErrors();
    TGraphErrors *graphBkg_2 = new TGraphErrors();
    TGraphErrors *graphSignificance_2 = new TGraphErrors();
    TGraphErrors *graphSig = new TGraphErrors();
    TGraphErrors *graphBkg = new TGraphErrors();
    TGraphErrors *graphSignificance = new TGraphErrors();
    // CMS [1 - 8.68], [10.09 - 12.86],  [14.18 - 19.00]
    double CMS_X[3]  = {4.84, 11.475, 16.59};
    double CMS_EX[3] = {3.84, 1.385,  2.41};
    double CMS_Y[3]  = {2.7, 4.1, 5.6};
    double CMS_EY[3] = {0.0, 0.0, 0.0};
    double CMS_sig[3]   = {22.1,25.9,45.1};
    double CMS_sig_E[3] = {8.1,6.3,8.0};
    double CMS_bkg[3]   = {49.0,14.0,20.0};
    double CMS_bkg_E[3] = {0.0,0.0,0.0};

    TGraphErrors *graphSignificanceCMS = new TGraphErrors(3,CMS_X,CMS_Y,CMS_EX,CMS_EY);
    TGraphErrors *graphSigCMS = new TGraphErrors(3,CMS_X,CMS_sig,CMS_EX,CMS_sig_E);
    TGraphErrors *graphBkgCMS = new TGraphErrors(3,CMS_X,CMS_bkg,CMS_EX,CMS_bkg_E);

    bool doFit = true;
    vector <double> binBoundaries = extraVar.isEquidistant ? extraVar.binEdgesEquidistant : extraVar.binEdges;
    for (int bin = 0; bin < nBins; bin++){
        YieldInfo YI_1 = yieldInQ2(1, getTMVAcut(1),bin,doFit);
        YieldInfo YI_2 = yieldInQ2(2, getTMVAcut(2),bin,doFit);
        YieldInfo YI_tmp = YI_1;
        YI_tmp.addYield(YI_2);
        binCenter = (binBoundaries.at(bin+1)+binBoundaries.at(bin))/2.0e6;
        binError  = (binBoundaries.at(bin+1)-binBoundaries.at(bin))/2.0e6;
        coutDebug("Bin center " + to_string(binCenter));
        coutDebug("Bin error " + to_string(binError));

        //Signal
        graphSig_1->SetPoint(graphSig_1->GetN(), binCenter,  YI_1.sigYield);
        graphSig_1->SetPointError(graphSig_1->GetN()-1, binError, YI_1.sigYieldErr);

        graphSig_2->SetPoint(graphSig_2->GetN(), binCenter,  YI_2.sigYield);
        graphSig_2->SetPointError(graphSig_2->GetN()-1, binError, YI_2.sigYieldErr);

        graphSig->SetPoint(graphSig->GetN(), binCenter,  YI_tmp.sigYield);
        graphSig->SetPointError(graphSig->GetN()-1, binError, YI_tmp.sigYieldErr);

        //Background
        graphBkg_1->SetPoint(graphBkg_1->GetN(), binCenter,  YI_1.bkgYield);
        graphBkg_1->SetPointError(graphBkg_1->GetN()-1, binError, YI_1.bkgYieldErr);

        graphBkg_2->SetPoint(graphBkg_2->GetN(), binCenter,  YI_2.bkgYield);
        graphBkg_2->SetPointError(graphBkg_2->GetN()-1, binError, YI_2.bkgYieldErr);

        graphBkg->SetPoint(graphBkg->GetN(), binCenter,  YI_tmp.bkgYield);
        graphBkg->SetPointError(graphBkg->GetN()-1, binError, YI_tmp.bkgYieldErr);

        //Significance
        graphSignificance_1->SetPoint(graphSignificance_1->GetN(), binCenter,  YI_1.significance);
        graphSignificance_1->SetPointError(graphSignificance_1->GetN()-1,binError,0);

        graphSignificance_2->SetPoint(graphSignificance_2->GetN(), binCenter,  YI_2.significance);
        graphSignificance_2->SetPointError(graphSignificance_2->GetN()-1,binError,0);

        graphSignificance->SetPoint(graphSignificance->GetN(), binCenter,  YI_tmp.significance);
        graphSignificance->SetPointError(graphSignificance->GetN()-1,binError,0);
    }

    design_YieldInQ2(1,  graphSig_1, graphBkg_1, graphSignificance_1, graphSignificanceCMS, fixRange);
    design_YieldInQ2(2,  graphSig_2, graphBkg_2, graphSignificance_2, graphSignificanceCMS, fixRange);
    design_YieldInQ2(12, graphSig,   graphBkg,   graphSignificance,   graphSignificanceCMS, fixRange);

    design_SignificanceInQ2(1,graphSignificance_1,graphSignificanceCMS,fixRange);
    design_SignificanceInQ2(2,graphSignificance_2,graphSignificanceCMS,fixRange);
    design_SignificanceInQ2(12,graphSignificance,graphSignificanceCMS,fixRange);
    return true;
}

bool ApplyCut(int Run, bool MC, bool Reference, bool PHSP, double TMVAcut){

    bool UseLowQ2Range = false;
    coutInfo("Staring to select signal events only for Run " + to_string(Run) + string(MC ? " MC" : " data") + string(Reference ? " Reference" : "") + string(PHSP? " PHSP" :"") + " at MVA reponse " + to_string(TMVAcut));


    bool onlyMuMu = MC && (!Reference && !PHSP);
    TChain *tree = get_BDT_TChain("2011",Run,MC,Reference && MC,PHSP,false,false);
    string cut = getFinalCut(MC, true, false,"", gammaTMdefault,
                             onlyMuMu, Reference, false, UseLowQ2Range,
                             false, TMefficiencyClass(), -1, TMVAcut, true);
    coutDebug("Cut: " + cut);

    string outputPath = GetFinalOutputFile(Run, MC, Reference, PHSP, false, UseLowQ2Range);
    coutDebug("Writting into new file " + outputPath);
    TFile *newFile = new TFile(outputPath.c_str(),"RECREATE");
    TTree *newTree = tree->CopyTree(cut.c_str(),"", tree->GetEntries());
    newFile->cd();
    newTree->Write("", TObject::kWriteDelete);
    newFile->Close();

    coutInfo("Done writting new file.");
    return 1;
}

bool ApplyCutPerYear(int Run, bool MC, bool Reference, bool PHSP, double TMVAcut){

    bool UseLowQ2Range = false;
    coutInfo("Staring to select signal events only for Run " + to_string(Run) + string(MC ? " MC" : " data") + string(Reference ? " Reference" : "") + string(PHSP? " PHSP" :"") + " at MVA reponse " + to_string(TMVAcut));

    bool onlyMuMu = MC && (!Reference && !PHSP);
    for(auto &year: yearsVector(MC,Reference,PHSP,Run)){
        TChain *tree = get_BDT_TChain(year,0,MC,Reference && MC,PHSP,false,false);
        string cut = getFinalCut(MC,true,false,"",gammaTMdefault,
                                 onlyMuMu,Reference,false,UseLowQ2Range,
                                 false,TMefficiencyClass(),-1,TMVAcut,true);
        coutDebug("Cut: " + cut);

        string outputPath = GetFinalOutputFile(year,Run, MC, Reference, PHSP, false, UseLowQ2Range);
        coutDebug("Writting into new file " + outputPath);
        TFile *newFile = new TFile(outputPath.c_str(),"RECREATE");
        TTree *newTree = tree->CopyTree(cut.c_str(),"", tree->GetEntries());
        newFile->cd();
        newTree->Write("", TObject::kWriteDelete);
        newFile->Close();

        coutInfo("Done writting new file.");
    }
    return 1;
}

bool ApplyCutToAll(int Run, double TMVAcut){
    if (!ApplyCut(Run, false, false, false, TMVAcut)) return 0;
    if (!ApplyCut(Run, true,  false, false, TMVAcut)) return 0;
    if (!ApplyCut(Run, true,  true,  false, TMVAcut)) return 0;
    if (!ApplyCut(Run, true,  false, true,  TMVAcut)) return 0;
    return 1;
}

bool ApplyCutPerYearAll(int Run, double TMVAcut){
    if (!ApplyCutPerYear(Run, false, false, false, TMVAcut)) return 0;
    if (!ApplyCutPerYear(Run, true,  false, false, TMVAcut)) return 0;
    if (!ApplyCutPerYear(Run, true,  true,  false, TMVAcut)) return 0;
    if (!ApplyCutPerYear(Run, true,  false, true,  TMVAcut)) return 0;
    return 1;
}
bool ApplyCutPerYearAll(int Run){
    ApplyCutPerYearAll(Run,getTMVAcut(Run));
    return 1;
}

bool printYileds(bool Rare){
    string filePath;    
    int entries = 0;
    //Print number of events passing stripping
    string sCut = getMuMucut();
    cout << "Trigger and online";
    for (auto& year : yearsData(12)){
        filePath = returnFileAddress(year,getRunID(year),"up",false,false,false,false,false,false);
        TChain *T = new TChain(treeName(false,false).c_str());
        T->Add(filePath.c_str());
        replace(filePath,"up","down");
        T->Add(filePath.c_str());
        if (Rare){
            TH1F *tmp1 = new TH1F("tmp1","tmp1",1002,-1.0,100.0);
            replace(sCut,"Q2","(J_psi_M**2)");
            T->Draw("J_psi_M >> tmp1", sCut.c_str());
            entries = tmp1->GetEntries();
        }
        else entries = T->GetEntries();
        cout << "\t& " << entries ;
        T->Clear();
    }
    cout << "\\\\" << endl;

    //Print number of events passing preselection
    cout << "Cut-based selection";
    for (auto& year : yearsData(12)){
        filePath = returnFileAddress(year,getRunID(year),"both",true,false,false,false,false,false);
        TFile *file = TFile::Open(filePath.c_str());
        TTree *tree = (TTree*)file->Get(treeName(false,true).c_str());
        if (Rare){
            TH1F *tmp2 = new TH1F("tmp2","tmp2",1002,-1.0,100.0);
            tree->Draw("Q2 >> tmp2", sCut.c_str());
            entries = tmp2->GetEntries();

        }
        else entries = tree->GetEntries();
        cout << "\t& " << entries ;
        tree->Clear();
        file->Close();
    }
    cout << "\\\\" << endl;

    //Print number of events passing BDT    for (auto& year : yearsData(12)){
    cout << "MVA selection";
    for (auto& year : yearsData(12)){
        filePath = returnFileAddress(year,getRunID(year),"both",true,true,false,false,false,false);
        TFile *file = TFile::Open(filePath.c_str());
        TTree *tree = (TTree*)file->Get(treeName(false,true).c_str());
        if (Rare){
            TH1F *tmp3 = new TH1F("tmp3","tmp3",1002,-1.0,100.0);
            tree->Draw("Q2 >> tmp3", sCut.c_str());
            entries = tmp3->GetEntries();
        }
        else entries = tree->GetEntries();
        cout << "\t& " << entries ;
        tree->Clear();
        file->Close();
    }
    cout << "\\\\" << endl;

    return true;
}