//Renata Kopecna


#include "GenLvlvsMC.hh"
#include <TH1D.h>
#include <TFile.h>
#include <TLegend.h>
#include <string>
#include <TGraphAsymmErrors.h>
#include <TGraphErrors.h>
#include <TGraph.h>
#include <design.hh>
#include <paths.hh>
#include <parse.hh>
#include <iostream>
#include <vector>
#include <helpers.hh>
#include "ScriptHelpers.hh"

std::string plotPath(std::string observable, std::string tag){
    return PLOTS_PATH+"MCfit/"+ "GenLvl_vs_MC_" + observable + tag +".eps";
}

std::vector<std::vector<double>> DavidsGenLvl = {
        { 0.5120, 0.1800, 0.1370, 0.1930, 0.2750, 0.4170, 0.4840, 0.5030}, //S1s
        { 0.2450, 0.6996, 0.7993, 0.7384, 0.6414, 0.4201, 0.3513, 0.3344}, //S1c
        { 0.1448, 0.0749, 0.0514, 0.0662, 0.0900, 0.1449, 0.1620, 0.1662}, //S2s
        {-0.1953,-0.6639,-0.7774,-0.7254,-0.6334,-0.4171,-0.3497,-0.3332}, //S2c
        { 0.0000, 0.0020,-0.0080,-0.0120,-0.0240,-0.0620,-0.1520,-0.2330}, //S3
        { 0.0870, 0.0040,-0.1080,-0.1950,-0.2510,-0.2770,-0.2900,-0.3040}, //S4
        { 0.2440, 0.0950,-0.1590,-0.3050,-0.4140,-0.4210,-0.3400,-0.2490}, //S5
        {-0.1270,-0.2070,-0.0860, 0.0950, 0.3020, 0.5140, 0.5570, 0.4700}, //S6s
        { 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000}, //S6c
        {-0.0060,-0.0100,-0.0030, 0.0020, 0.0030,-0.0020,-0.0020, 0.0050}, //S7
        { 0.0010,-0.0030,-0.0040,-0.0030,-0.0010, 0.0030,-0.0010, 0.0030}, //S8
        {-0.0020, 0.0000, 0.0020, 0.0040,-0.0020,-0.0010,-0.0010, 0.0010}  //S9
    }; //I would kill for a dictionary in C++ now

double getDavidsValue(std::string observable, int bin){ //Not the most efficienc implementation, but oh well, hard to live without numpy
    if (observable == "Fl")  return (1-4.0/3.0*DavidsGenLvl[0][bin]);
    if (observable == "S1s") return DavidsGenLvl[0][bin];
    if (observable == "S1c") return DavidsGenLvl[1][bin];
    if (observable == "S2s") return DavidsGenLvl[2][bin];
    if (observable == "S2c") return DavidsGenLvl[3][bin];
    if (observable == "S3")  return DavidsGenLvl[4][bin];
    if (observable == "S4")  return DavidsGenLvl[5][bin];
    if (observable == "S5")  return DavidsGenLvl[6][bin];
    if (observable == "S6s") return DavidsGenLvl[7][bin];
    if (observable == "Afb") return 3.0/4.0*DavidsGenLvl[7][bin];
    if (observable == "S6c") return DavidsGenLvl[8][bin];
    if (observable == "S7")  return DavidsGenLvl[9][bin];
    if (observable == "S8")  return DavidsGenLvl[10][bin];
    if (observable == "S9")  return DavidsGenLvl[11][bin];

    spdlog::critical("I don't know observable "+observable);
    spdlog::critical("Throwing an error now.");
    assert(0);
    return 0;
}

TGraphAsymmErrors* DavidsGenLvlGraph(std::string observable){
    int nBins = DavidsGenLvl[0].size(); //In case the vector above changes, this should be secure
    std::vector<double> q2min = get_TheQ2binsmin(nBins,false);
    std::vector<double> q2max = get_TheQ2binsmax(nBins,false);

    //Initilialize the TGraph
    //Initialize it on purpose with 0 points; in case something goes wrong with nBins, this makes sure there are no segfaults
    TGraphAsymmErrors* graph = new TGraphAsymmErrors(0);

    //Loop over bins to fill TGraph from DavidsGenLvl
    for (int b = 0; b < nBins; b++){
        graph->SetPoint(graph->GetN(),bin_center_q2(q2min,q2max,b), getDavidsValue(observable,b));
        graph->SetPointError(graph->GetN()-1,bin_halfWidth_q2(q2min,q2max,b), bin_halfWidth_q2(q2min,q2max,b), 0, 0);
    }
    return graph;

}

int compareGenLvlMC(std::vector<TGraphAsymmErrors*> graphs,
                    std::vector<std::string> legends,
                    std::vector<int> color,
                    std::vector<int> markerStyle,
                    std::string observable, std::string plotPath,
                    std::vector<int> plotDiff){ //What two plots should be considered for plotting the "pulls"? Marked with 1 and 2, plotting 1-2

    int idx_1 = find(plotDiff.begin(), plotDiff.end(), 1)-plotDiff.begin();
    int idx_2 = find(plotDiff.begin(), plotDiff.end(), 2)-plotDiff.begin();
    bool plotD = (idx_1 != int(graphs.size()) && idx_2 != int(graphs.size()));
    spdlog::debug("plot diff? " + boolToString(plotD));
    spdlog::debug("Idx_1: {0:d}\tIdx_2: {1:d}", idx_1, idx_2);

    //Create a canvas
    TCanvas *c_compare = new TCanvas("c_compare", "c_compare", 1600, 1200);

    const double pullHeight = 0.27;
    const double pullFrameRange = 3.5;
    TPad *pad1 = new TPad("pad1", "plot",0.0,pullHeight,1.0,1.0,0);
    TPad *pad2 = new TPad("pad2", "pull",0.0,0.0,1.0,pullHeight,0);
    if(plotD){
        pad1->Draw();
        pad2->Draw();
        pad1->SetBorderSize  (0);
        pad1->SetMargin(0.125,0.05,1e-6,1.25/(1.0 - pullHeight));
        pad1->SetTickx();
        pad1->SetTicky();
        pad1->cd();
    }
    else{
        c_compare->cd();
        c_compare->SetMargin(0.125,0.05,0.1,0.1);
        c_compare->SetTickx();
        c_compare->SetTicky();
    }


    //Create an empty hist for the axes    
    TH1D* haxes = emptyHist(observable);
    haxes->Draw("AXIS");

    //Draw boxes for resonances
    drawResonances(plotD ? pad1 : c_compare, obsv_range(observable)[0], obsv_range(observable)[1]);

    haxes->SetFillStyle(0);
    haxes->Draw("AXIS SAME");

    TLegend* leg = new TLegend(0.62,0.70,0.94,0.94);
    leg->SetBorderSize(0);
    leg->SetTextFont(132);
    leg->SetFillColor(0);

    for_indexed(auto g: graphs){
        designTGraph(g, markerStyle[i], color[i]);
        if (legends[i]!="") leg->AddEntry(g, legends[i].c_str());
        g->Draw("P");
    }
    leg->Draw("SAME");

    TH1D *diff = nullptr;
    if (plotD){ //If there are no tags 1 and 2, don't do anything
        pad2->Clear();
        pad2->SetMargin(0.125,0.05,0.1/ pullHeight, 1e-6);
        pad2->SetTickx();
        pad2->cd();

        std::string y_title = "Diff [#sigma]";//legends[idx_1] + "-" + legends[idx_2];
        diff = hist_graphDiff(graphs[idx_1],graphs[idx_2],pullFrameRange, y_title);

        double max = haxes->GetBinLowEdge(haxes->GetNbinsX()+1);
        double min = haxes->GetBinLowEdge(1);
        //I hate that I have to workaround ranges of an efin histogram

        TLine *sigmaUp_3  = threeSigmaLine(min,max, true);
        TLine *sigmaLow_3 = threeSigmaLine(min,max, false);
        TLine *sigmaUp_1  = oneSigmaLine(min,max, true);
        TLine *sigmaLow_1 = oneSigmaLine(min,max, false);
        design_pull(diff, 9011, 9011, pullHeight+0.1, pullFrameRange);
        diff->GetYaxis()->SetNdivisions(110);
        diff->Draw("");

        sigmaUp_3->Draw("same");
        sigmaLow_3->Draw("same");
        sigmaUp_1->Draw("same");
        sigmaLow_1->Draw("same");
    }

    spdlog::info("Saving into "+ plotPath);
    c_compare->Print(plotPath.c_str(), "eps");
    if (spdlog_debug()){ //When debugging save also the root file
        replace(plotPath,".eps",".root");
        c_compare->Print(plotPath.c_str(), "root");
    }

    delete c_compare;
    delete haxes;
    delete diff;
    return 0;

}

int plotComparison_Toy_GenVsMC(std::string observable, bool isRef){
    bool onlySig = false;
    bool onlyBkg = false;
    bool onlyAngles = false;
    bool onlyMass = false;
    bool bkgFromLowMass = false;
    bool bkgFromHighMass = true;

    basic_params params_MC = basic_params();
    params_MC.Run = 12;
    params_MC.nBins = isRef ? 1 : 4;
    params_MC.reference = isRef;

    basic_params params_Toys = basic_params();
    params_Toys.Run = 12;
    params_Toys.nBins = isRef ? 1 : 4;
    params_Toys.reference = isRef;

    std::vector<std::string> legends;
    std::vector<int> color;
    std::vector<int> markerStyle;
    std::vector<TGraphAsymmErrors*> graphs;
    std::vector<int> difference; //makes a plot with difference of grpahs with a tag of 1-2
    //I shoudl write a class for this, but too much effort


    //Add init parameters
    std::string path = final_result_name_MC(params_MC, params_Toys.nBins, isRef, false, true, false, false);
    graphs.push_back(get_TGraphFromFile(path,observable,isRef));
    legends.push_back("MC"); //Don't plot legend as it is same as sigMC
    color.push_back(9012);
    markerStyle.push_back(27);
    difference.push_back(1);

    //Fill the toy init parameters
    path = init_params_name_toys(1,isRef,params_Toys.nBins, true, params_Toys, params_Toys.Run,  onlyAngles, onlyMass, onlySig, onlyBkg, bkgFromLowMass, bkgFromHighMass);
    graphs.push_back(get_TGraphFromFile(path,observable,isRef));
    legends.push_back("Toy");
    color.push_back(9007);
    markerStyle.push_back(34);
    difference.push_back(2);

    std::string plotPath = PLOTS_PATH+"Toys/"+ "MC_vs_Init_" + observable +".eps";
    return compareGenLvlMC(graphs,legends,color,markerStyle,observable, plotPath, difference);

}

//This one compares the data-like fit to the init toy values
int plotComparison_Toy_InitVsAcutalFit(std::string observable, bool isRef){

    bool onlySig = false;
    bool onlyBkg = false;
    bool onlyAngles = false;
    bool onlyMass = false;
    bool bkgFromLowMass = false;
    bool bkgFromHighMass = false;

    double fractionOfStats = 1.0;

    basic_params params_Init= basic_params();
    params_Init.Run = 12;
    params_Init.nBins = isRef ? 1 : 4;
    params_Init.reference = isRef;
    params_Init.jobID = -1;

    basic_params params_Fit= basic_params();
    params_Fit.Run = 12;
    params_Fit.nBins = isRef ? 1 : 4;
    params_Fit.reference = isRef;

    std::vector<std::string> legends;
    std::vector<int> color;
    std::vector<int> markerStyle;
    std::vector<TGraphAsymmErrors*> graphs;
    std::vector<int> difference; //makes a plot with difference of grpahs with a tag of 1-2
    //I shoudl write a class for this, but too much effort

    //Add init parameters
    std::string paramFile = init_params_name_toys(-1,params_Init.reference, params_Init.nBins, true, params_Init, params_Init.Run, onlyAngles, onlyMass, onlySig, onlyBkg, bkgFromLowMass, bkgFromHighMass);
    graphs.push_back(get_TGraphFromFile(paramFile,observable,isRef));
    legends.push_back("Init"); //Don't plot legend as it is same as sigMC
    color.push_back(9012);
    markerStyle.push_back(27);
    difference.push_back(1);

    //Add fitted parameters
    std::string fitFile = final_result_name(isRef, false, params_Fit,
                                            true, params_Fit.nBins, params_Fit.Run,
                                            true, fractionOfStats,
                                            false, -1);

    graphs.push_back(get_TGraphFromFile(fitFile,observable,isRef));
    legends.push_back("Fit");
    color.push_back(9007);
    markerStyle.push_back(34);
    difference.push_back(2);

    std::string plotPath = PLOTS_PATH+"Toys/"+ "FinalToy_Init_vs_Fit_" + observable + std::string(isRef ? "_Ref" : "" )+".eps";
    return compareGenLvlMC(graphs,legends,color,markerStyle,observable, plotPath, difference);
}


int plotComparison_Toy_GenVsFit(std::string observable, bool isRef){

    bool onlySig = false;
    bool onlyBkg = false;
    bool onlyAngles = false;
    bool onlyMass = false;
    bool bkgFromLowMass = false;
    bool bkgFromHighMass = true;

    basic_params params_Init= basic_params();
    params_Init.Run = 12;
    params_Init.nBins = isRef ? 1 : 4;
    params_Init.reference = isRef;

    basic_params params_Fit= basic_params();
    params_Fit.Run = 12;
    params_Fit.nBins = isRef ? 1 : 4;
    params_Fit.reference = isRef;

    std::vector<std::string> legends;
    std::vector<int> color;
    std::vector<int> markerStyle;
    std::vector<TGraphAsymmErrors*> graphs;
    std::vector<int> difference; //makes a plot with difference of grpahs with a tag of 1-2
    //I shoudl write a class for this, but too much effort

    //Add init parameters
    std::string path = init_params_name_toys(1,isRef,params_Init.nBins, true, params_Init, params_Init.Run,  onlyAngles, onlyMass, onlySig, onlyBkg, bkgFromLowMass, bkgFromHighMass);
    graphs.push_back(get_TGraphFromFile(path,observable,isRef));
    legends.push_back("Init"); //Don't plot legend as it is same as sigMC
    color.push_back(9012);
    markerStyle.push_back(27);
    difference.push_back(1);

    //Add fitted parameters
    path = final_result_name_toys(1,isRef,params_Fit.nBins, true, params_Fit, params_Fit.Run,  onlyAngles, onlyMass, onlySig, onlyBkg, bkgFromLowMass, bkgFromHighMass);
    graphs.push_back(get_TGraphFromFile(path,observable,isRef));
    legends.push_back("Fit");
    color.push_back(9007);
    markerStyle.push_back(34);
    difference.push_back(2);

    std::string plotPath = PLOTS_PATH+"Toys/"+ "Init_vs_Fit_" + observable + std::string(isRef ? "_Ref" : "" )+".eps";
    return compareGenLvlMC(graphs,legends,color,markerStyle,observable, plotPath, difference);
}

int plotComparisonMCFolding(std::string observable){
    basic_params params_MC= basic_params();
    params_MC.Run = 12;
    params_MC.nBins = 4;
    params_MC.reference = false;

    std::vector<std::string> legends;
    std::vector<int> color;
    std::vector<int> markerStyle;
    std::vector<TGraphAsymmErrors*> graphs;
    std::vector<int> difference; //makes a plot with difference of grpahs with a tag of 1-2

    for (int f = -1; f <5; f++){
        params_MC.folding = f;
        if (f == -1)  difference.push_back(1);
        if (plotObsv(f,observable)){
            graphs.push_back(get_TGraphFromFile(final_result_name_MC(params_MC,params_MC.nBins, false, false, true, false, false),observable,false));
            legends.push_back("Fld " + std::to_string(f));
            color.push_back(9009+f);
            markerStyle.push_back(25+f);
            difference.push_back(2);
        }
    }

    std::string tag = "";
    std::string plotPath = PLOTS_PATH+"MCfit/"+ "Folding_" + observable + tag +".eps";

    assert(!compareGenLvlMC(graphs,legends,color,markerStyle,observable, plotPath, difference));

    return 0;
}

int plotComparison(std::string observable){

    basic_params params_genLvl = basic_params();
    params_genLvl.Run = 2;
    params_genLvl.year = 2017;
    params_genLvl.nBins = 8;

    bool PHSP = false;

    std::vector<std::string> legends;
    std::vector<int> color;
    std::vector<int> markerStyle;
    std::vector<TGraphAsymmErrors*> graphs;
    std::vector<int> difference; //makes a plot with difference of grpahs with a tag of 1-2
    //I shoudl write a class for this, but too much effort

    //First make a comparison between my and David's genLvl
    //---------------------------------------------------//    
    //Add genLvl
    graphs.push_back(get_TGraphFromFile(final_result_name_genLvlMC(params_genLvl, 8, PHSP, true, false),observable,false));
    legends.push_back("GenLvl MC"); //Don't plot legend as it is same as sigMC
    color.push_back(9012);
    markerStyle.push_back(27);
    difference.push_back(1);

    //Fill David's results first
    graphs.push_back(DavidsGenLvlGraph(observable));
    legends.push_back("David's GenLvl results");
    color.push_back(9007);
    markerStyle.push_back(34);
    difference.push_back(2);

    std::string tag = "_MyVsDavids";
    std::string plotPath = PLOTS_PATH+"MCfit/"+ "GenLvl_vs_MC_" + observable + tag +".eps";

    assert(!compareGenLvlMC(graphs,legends,color,markerStyle,observable, plotPath, difference));

    //Delete David's graphs
    graphs.clear();
    legends.clear();
    color.clear();
    markerStyle.clear();
    difference.clear();
    //---------------------------------------------------//
    //Compare my fit of genLvl mc and processed MC
    //---------------------------------------------------//
    //Add MC as it should be done

    basic_params params_MC= basic_params();
    params_MC.Run = 12;
    params_MC.nBins = 4;
    params_MC.reference = false;

    graphs.push_back(get_TGraphFromFile(final_result_name_MC(params_MC, 4, false, PHSP, true, false, false),observable,false));
    legends.push_back("LHCb simulation");
    color.push_back(9016);
    markerStyle.push_back(24);    
    difference.push_back(1);

    //Add genLvl MC
    params_genLvl.nBins = 4;
    graphs.push_back(get_TGraphFromFile(final_result_name_genLvlMC(params_genLvl, 4, PHSP, true, false),observable,false));
    legends.push_back("Generator level");
    color.push_back(9010);
    markerStyle.push_back(25);
    difference.push_back(2);

    plotPath = PLOTS_PATH+"MCfit/"+ "GenLvl_vs_MC_" + observable + ".eps";
    return compareGenLvlMC(graphs,legends,color,markerStyle,observable, plotPath, difference);
    //---------------------------------------------------//

}