//Renata Kopecna


#include "ScriptHelpers.hh"
#include <paths.hh>
#include <TFile.h>
#include <TTree.h>
#include <helpers.hh>
#include <design.hh>
#include <constants.hh>

#include <spdlog.h>

//Takes two valErr and returns the difference in sigmas
double sigmaAway(valErr a, valErr b){
    double diff = a.val - b.val;
    double effErr = sqrt(pow(a.err,2)+pow(b.err,2));
    return effErr == 0 ? 0 : diff/effErr;
}

//Returns true if the observable is measured in certain folding
bool plotObsv(int folding, std::string obs){
    if (obs == "Fl" || obs == "S3") return true; //FL and S3 always there

    switch(folding){
    case -1: return true;
    case 0:
        if (obs == "Afb") return true;
        else if (obs == "S9") return true;
        else return false;
    case 1:
        if (obs == "S4") return true;
        else return false;
    case 2:
        if (obs == "S5") return true;
        else return false;
    case 3:
        if (obs == "S7") return true;
        else return false;
    case 4:
        if (obs == "S8") return true;
        else return false;
    }
    return false;
}

std::vector<double> obsv_range(std::string parName){
    if (parName == "S1s")       return {+0.0, +1.0};
    if (parName == "Fl")        return {+0.0, +1.0};
    if (parName == "S3")        return {-0.5, +0.3};
    if (parName == "S4")        return {-0.5, +0.3};
    if (parName == "S5")        return {-0.5, +0.3};
    if (parName == "Afb")       return {-0.2, +1.0};
    if (parName == "S7")        return {-0.25, +0.25};
    if (parName == "S8")        return {-0.25, +0.30};
    if (parName == "S9")        return {-0.25, +0.25};

    //mass fit
    if (parName == "m_b")       return {PDGMASS_B-10,PDGMASS_B+10};
    if (parName == "m_sigma_1") return {20.0, 40.0};
    if (parName == "alpha_1")   return {0.5, 2.5};
    if (parName == "alpha_2")   return {0.5, 1.5};
    if (parName == "n_1")       return {0.1, 15.0};
    if (parName == "n_2")       return {5.0, 15.0};
    if (parName == "m_lambda")  return {-0.1, -0.00001};
    if (parName == "f_sig")     return {0.0, 1.0};


    //angular bkg
    if (parName == "cbkgctl0")  return {-10.0, 10.0};
    if (parName == "cbkgctl1")  return {-10.0, 10.0};
    if (parName == "cbkgctl2")  return {-10.0, 10.0};
    if (parName == "cbkgctl3")  return {-10.0, 10.0};
    if (parName == "cbkgctl4")  return {-10.0, 10.0};
    if (parName == "cbkgctk0")  return {-10.0, 10.0};
    if (parName == "cbkgctk1")  return {-10.0, 10.0};
    if (parName == "cbkgctk2")  return {-10.0, 10.0};
    if (parName == "cbkgctk3")  return {-10.0, 10.0};
    if (parName == "cbkgctk4")  return {-10.0, 10.0};
    if (parName == "cbkgctk5")  return {-10.0, 10.0};
    if (parName == "cbkgctk6")  return {-10.0, 10.0};
    if (parName == "cbkgphi0")  return {-10.0, 10.0};
    if (parName == "cbkgphi1")  return {-10.0, 10.0};
    if (parName == "cbkgphi2")  return {-10.0, 10.0};
    if (parName == "cbkgphi3")  return {-10.0, 10.0};
    if (parName == "cbkgphi4")  return {-10.0, 10.0};
    if (parName == "cbkgmkpi0")  return {-10.0, 10.0};
    if (parName == "cbkgmkpi1")  return {-10.0, 10.0};
    if (parName == "cbkgmkpi2")  return {-10.0, 10.0};
    if (parName == "cbkgmkpi3")  return {-10.0, 10.0};
    if (parName == "cbkgmkpi4")  return {-10.0, 10.0};

    spdlog::critical("Wrong parameter tree name! Setting range to default  -1.0, 1.0");
    return {-1.0,1.0};
}


void designTGraph(TGraphAsymmErrors *g, int markerStyle, int color){
    g->SetLineColor(color);
    g->SetMarkerColor(color);
    g->SetLineWidth(2);
    g->SetMarkerStyle(markerStyle);
    g->SetMarkerSize(3);
    g->SetFillColor(0);
    g->SetFillStyle(0);
}
TH1D *emptyHist(std::string observable, double min, double max, int nBins){
    TH1D* haxes = new TH1D("haxes", std::string(";#it{q}^{2} ["+GeVc_2()+"];"+observable).c_str(), nBins, Q2_MIN_RANGE, Q2_MAX_RANGE_B0);
    haxes->SetMinimum(min);
    haxes->SetMaximum(max);
    haxes->SetTitleSize(0.06, "XY");
    haxes->SetLabelSize(0.05, "XY");
    return haxes;
}

TH1D *emptyHist(std::string observable){
    double min = obsv_range(observable)[0]; //Set the range for each parameter
    double max = obsv_range(observable)[1]; //Needs to run twice but whatever
    return (emptyHist(observable,min,max,100));
}


TGraphAsymmErrors* get_TGraphFromFile(std::string fileName, std::string observable, bool isRef){
    //Open file
    spdlog::info("Opening " + fileName);
    TFile* file = new TFile(fileName.c_str(), "READ");
    if (!file->GetListOfKeys()->Contains(observable.c_str())){
        spdlog::critical("Wrong tree name " + observable + "! Abort.");
        assert(0);
    }

    TTree* tree = (TTree*)file->Get(observable.c_str());
    tree->SetBranchStatus("*",1); //Activate all branches

    //Read the q2 bins
    int totBins = DEFAULT_TREE_VAL;
    tree->SetBranchAddress("totBins", &totBins);
    tree->GetEntry(0);

    std::vector<double> q2min = get_TheQ2binsmin(totBins,isRef);
    std::vector<double> q2max = get_TheQ2binsmax(totBins,isRef);

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

    //Loop over nBins that are saved in the tree
    for (int b = 0; b < totBins; b++){
        double value    = DEFAULT_TREE_VAL;
        double error    = DEFAULT_TREE_ERR;
        double errorup  = DEFAULT_TREE_ERR;
        double errordown= DEFAULT_TREE_ERR;
        int migrad  = DEFAULT_TREE_INT;
        int cov     = DEFAULT_TREE_INT;

        tree->SetBranchAddress("value", &value);
        tree->SetBranchAddress("error", &error);
        tree->SetBranchAddress("error_up", &errorup);
        tree->SetBranchAddress("error_down", &errordown);
        tree->SetBranchAddress("migrad", &migrad);
        tree->SetBranchAddress("status_cov", &cov);
        tree->GetEntry(b);

        //if opts.hesse is non-active, use error for both error_up and error_down
        if(errordown == DEFAULT_TREE_ERR) errordown = error;
        if(errorup   == DEFAULT_TREE_ERR) errorup   = error;

        //Set point in TGraph
        spdlog::debug("Setting point in bin {0:d} in Tgraph {1:s}", b, graph->GetName());
        spdlog::debug("bin center {0:f}, value {1:f}", bin_center_q2(q2min,q2max,b), value);
        spdlog::debug("bin width {0:f}, err up {1:f}, err down {2:f}",
                      bin_halfWidth_q2(q2min,q2max,b), fabs(errordown), fabs(errorup));
        graph->SetPoint(graph->GetN(),bin_center_q2(q2min,q2max,b), value);
        graph->SetPointError(graph->GetN()-1,bin_halfWidth_q2(q2min,q2max,b), bin_halfWidth_q2(q2min,q2max,b), fabs(errordown), fabs(errorup));

    } //loop over bins    

    spdlog::debug("Graph filled.");
    file->Close();

    return graph;
}


TH1D* hist_graphDiff(TGraphAsymmErrors* graph_a, TGraphAsymmErrors* graph_b,
                     double range, std::string title){
    //Check they have the same number of points
    int n_points = graph_a->GetN();
    if (n_points!= graph_b->GetN()){
        spdlog::error("The input graphs have different number of points!");
        spdlog::error("Graph a: {0:d}\tGraph b: {1:d}", n_points, graph_b->GetN());
        assert(0);
    }
    //Technically one shoudl check matching x-axis values, but screw that

    TH1D* hist = emptyHist(title,-range,range,1000);
    //Make a histogram with the difference

    for (int p = 0; p<n_points; p++){
        double val = graph_a->GetY()[p] - graph_b->GetY()[p];
        double effErr = 0;
        if (val > 0) effErr = sqrt(pow(graph_a->GetErrorYlow(p),2) +pow(graph_b->GetErrorYhigh(p),2));
        else         effErr = sqrt(pow(graph_a->GetErrorYhigh(p),2)+pow(graph_b->GetErrorYlow(p),2));
        double low_edge  = hist->GetXaxis()->FindBin(graph_a->GetX()[p]-graph_a->GetErrorXlow(p));
        double high_edge = hist->GetXaxis()->FindBin(graph_a->GetX()[p]+graph_a->GetErrorXhigh(p));
        spdlog::debug("low_edge {0:d}",  low_edge);
        spdlog::debug("high_edge {0:d}", high_edge);
        for (int b = low_edge+1; b<high_edge; b++){
            hist->SetBinContent(b,effErr == 0 ? 0 : val/effErr);
        }
    }

    return hist;
}