#include "apply_plot_style.h"

void ApplyPlotStyle::test(){

    string datafile = "~/Documents/Work/MICE Data/Data run 7469/MAUS 3pt0/data.root";

    TFile* data = new TFile(datafile.c_str());

    TTree* T = (TTree*)data->Get("T");

    TH1F* x = new TH1F("x", ";x (mm)", 150, -150.0, 150.0);

    TH1F* y = new TH1F("y", ";y", 150, -150.0, 150.0);

    TH2F* xy_all = new TH2F("xy_all", ";x (mm); y(mm)", 150, -150.0, 150.0, 150, -150.0, 150.0);

    TH2F* xy_some = new TH2F("xy_some", ";x (mm); y(mm)", 150, -150.0, 150.0, 150, -150.0, 150.0);

    TRandom3 random;

    double a, b;

    T->SetBranchAddress("TKU_s1_x", &a);

    T->SetBranchAddress("TKU_s1_y", &b);

    double cut_x_min = -50.0;

    double cut_x_max = 50.0;

    double cut_y_min = -50.0;

    double cut_y_max = 50.0;

    for(int i = 0; i < T->GetEntries(); i++){

        T->GetEntry(i);

        if(b > cut_y_min && b < cut_y_max){

            if(a > cut_x_min && b < cut_x_max){

                x->Fill(a);

                y->Fill(b);

                xy_all->Fill(a, b);

                if(random.Uniform() < 0.7){

                    xy_some->Fill(a, b);

                }

        }

        }

    }

    Hist1D(x, y, "style_test.pdf", -50.0, 50.0);

    Hist2D(xy_all, xy_some, "style_test_2d.pdf", -50.0, 50.0, -50.0, 50.0, 0);

}

ApplyPlotStyle::ApplyPlotStyle(){

    std::cout << "Plot style object created.\n";

}

void ApplyPlotStyle::Hist1D(TH1F* hist1, TH1F* hist2, string saveAs, double cut_lower, double cut_upper, bool normalise, string legend1, string legend2){

    /*

     * Stack data and MC hists together onto one plot:

     *   - Normalise to 1

     *   - Set plot style

     *   - Stack

     */

    hist1->Sumw2(); // makes statistical error bars

    double norm_hist1 = hist1->Integral();

    if(normalise == true)

        hist1->Scale(1.0/norm_hist1);

    double norm_hist2 = hist2->Integral();

    if(normalise == true)

        hist2->Scale(1.0/norm_hist2);

    std::cout << "Hist 1 entries: " << norm_hist1 << ", Hist 2 entries: "<< norm_hist2 << "\n";

    hist1->SetMarkerStyle(20);

    hist1->SetMarkerSize(0.5);

    hist1->SetMarkerColor(kBlack);

    hist1->SetLineColor(kBlack);

    hist2->SetLineStyle(1);

    hist2->SetLineWidth(1);

    hist2->SetLineColor(kBlack);

    hist2->SetFillColor(kOrange-2);

    THStack* hist_stack = new THStack("hist_stack", "some_title");

    hist_stack->Add(hist2, "hist");

    hist_stack->Add(hist1, "e1");

    string axis_title = ";" + string(hist1->GetXaxis()->GetTitle());

    hist_stack->SetTitle(axis_title.c_str());

    TLegend legend(0.65, 0.75, 0.89, 0.89);

    legend.AddEntry(hist1, legend1.c_str(), "ELP");

    legend.AddEntry(hist2, legend2.c_str(), "LF");

    TCanvas c("c", "c");

    c.SetBorderSize(0);

    hist_stack->Draw("nostack");

    hist_stack->GetXaxis()->SetTitleSize(.05);

    hist_stack->GetXaxis()->SetLabelSize(.05);

    hist_stack->GetYaxis()->SetLabelSize(.05);

    c.Update();

    double y_min = c.GetUymin();

    double y_max = c.GetUymax();

    // Now sort out any cut-indicators:

  //  if(cut_lower != TMath::Infinity()){

        // make a lower-boundary line

        TLine lower_line(cut_lower, y_min, cut_lower, y_max);

        lower_line.SetLineColor(kBlack);

        lower_line.SetLineWidth(3);

        lower_line.Draw();

  //  }

  //  if(cut_upper != TMath::Infinity()){

        // make a lower-boundary line

        TLine upper_line(cut_upper, y_min, cut_upper, y_max);

        upper_line.SetLineColor(kBlack);

        upper_line.SetLineWidth(3);

        upper_line.Draw();

  //  }

  //  c.Update();

    legend.SetTextSize(.05);

    legend.Draw();

    c.Print(saveAs.c_str());

}

void ApplyPlotStyle::Hist1D_single(TH1F* hist, string saveAs="output.pdf", double cut_lower=TMath::Infinity(), double cut_upper=TMath::Infinity(), string legend1="Data", bool isData=true){

    hist->Sumw2(); // makes statistical error bars

    if(isData == true){

        hist->SetMarkerStyle(20);

        hist->SetMarkerSize(0.5);

        hist->SetMarkerColor(kBlack);

        hist->SetLineColor(kBlack);

    }

    else{

        hist->SetLineStyle(1);

        hist->SetLineWidth(1);

        hist->SetLineColor(kBlack);

        hist->SetFillColor(kOrange-2);

    }

    TLegend legend(0.65, 0.75, 0.89, 0.89);

    legend.AddEntry(hist, legend1.c_str(), "ELP");

    TCanvas c("c", "c");

    c.SetBorderSize(0);

    hist->Draw("hist");

    hist->GetXaxis()->SetTitleSize(.05);

    hist->GetXaxis()->SetLabelSize(.05);

    hist->GetYaxis()->SetLabelSize(.05);

    c.Update();

    double y_min = c.GetUymin();

    double y_max = c.GetUymax();

    // Now sort out any cut-indicators:

    //  if(cut_lower != TMath::Infinity()){

    // make a lower-boundary line

    TLine lower_line(cut_lower, y_min, cut_lower, y_max);

    lower_line.SetLineColor(kBlack);

    lower_line.SetLineWidth(3);

    lower_line.Draw();

    //  }

    //  if(cut_upper != TMath::Infinity()){

    // make a lower-boundary line

    TLine upper_line(cut_upper, y_min, cut_upper, y_max);

    upper_line.SetLineColor(kBlack);

    upper_line.SetLineWidth(3);

    upper_line.Draw();

    //  }

    //  c.Update();

    legend.SetTextSize(.05);

    legend.Draw();

    c.Print(saveAs.c_str());

}

void ApplyPlotStyle::Hist2D(TH2F* data_hist, TH2F* mc_hist, string saveAs, double cut_a_lower, double cut_a_upper, double cut_b_lower, double cut_b_upper, int legend_location, bool normalise, string legend1, string legend2){

    /*

     * Stack data and MC hists together onto one plot:

     *   - Normalise to 1

     *   - Set plot style

     *   - Stack

     * Legend (stats box) location:

     *   0: no legend

     *   1: top left

     *   2: top right

     *   3: bottom left

     *   4: bottom right

     */

    double norm_data = data_hist->GetMaximum();

    std::cout << "Max in data hist = " << norm_data << "\n";

    double norm_mc = mc_hist->GetMaximum();

    std::cout << "Max in mc hist = " << norm_mc << "\n";

    if(normalise == true){

        if(norm_data > norm_mc){

            data_hist->SetMinimum(0.0);

            data_hist->SetMaximum(norm_data);

            mc_hist->SetMinimum(0.0);

            mc_hist->SetMaximum(norm_data);

        }

        else{

            data_hist->SetMinimum(0.0);

            data_hist->SetMaximum(norm_mc);

            mc_hist->SetMinimum(0.0);

            mc_hist->SetMaximum(norm_mc);

        }

    }

    data_hist->GetYaxis()->SetTitleOffset(1.5);

    mc_hist->GetYaxis()->SetTitleOffset(1.5);

    data_hist->GetXaxis()->SetTitleSize(.05);

    data_hist->GetYaxis()->SetTitleSize(.05);

    data_hist->GetXaxis()->SetLabelSize(.05);

    data_hist->GetYaxis()->SetLabelSize(.05);

    data_hist->GetZaxis()->SetLabelSize(.05);

    mc_hist->GetXaxis()->SetTitleSize(.05);

    mc_hist->GetYaxis()->SetTitleSize(.05);

    mc_hist->GetXaxis()->SetLabelSize(.05);

    mc_hist->GetYaxis()->SetLabelSize(.05);

    mc_hist->GetZaxis()->SetLabelSize(.05);

    data_hist->SetTitle(legend1.c_str());

    mc_hist->SetTitle(legend2.c_str());

    TCanvas c("c", "c", 800, 400);

    c.Divide(2, 1);

    c.SetFillColor(0);

    c.SetFrameFillStyle(0);

    c.cd(1);

    data_hist->Draw("colz");

    c.cd(2);

    mc_hist->Draw("colz");

    if(legend_location == 1){

        gStyle->SetStatX(0.33);

        gStyle->SetStatY(0.88);

        gStyle->SetOptStat("emr");

    }

    else if(legend_location == 2){

        gStyle->SetStatX(0.88);

        gStyle->SetStatY(0.88);

        gStyle->SetOptStat("emr");

    }

    else if(legend_location == 3){

        gStyle->SetStatX(0.33);

        gStyle->SetStatY(0.38);

        gStyle->SetOptStat("emr");

    }

    else if(legend_location == 4){

        gStyle->SetStatX(0.88);

        gStyle->SetStatY(0.38);

        gStyle->SetOptStat("emr");

    }

    else{

        gStyle->SetOptStat(0);

    }

    c.cd(1);

    gPad->Update();

    /*

    TPaletteAxis *palette = (TPaletteAxis*)data_hist->GetListOfFunctions()->FindObject("palette");

    // the following lines moe the paletter. Choose the values you need for the position.

    palette->SetX1NDC(0.9);

    palette->SetX2NDC(0.92);

    palette->SetY1NDC(0.1);

    palette->SetY2NDC(0.9);

    gPad->Modified();

    gPad->Update();

     */

    double x_min = gPad->GetUxmin();

    double x_max = gPad->GetUxmax();

    double y_min = gPad->GetUymin();

    double y_max = gPad->GetUymax();

    // Now sort out any cut-indicators:

//    if(cut_a_lower != TMath::Infinity()){

        // make a lower-boundary line

        TLine lower_a_line(cut_a_lower, y_min, cut_a_lower, y_max);

        lower_a_line.SetLineColor(kBlack);

        lower_a_line.SetLineWidth(3);

        lower_a_line.Draw();

//    }

//    if(cut_a_upper != TMath::Infinity()){

        // make a lower-boundary line

        TLine upper_a_line(cut_a_upper, y_min, cut_a_upper, y_max);

        upper_a_line.SetLineColor(kBlack);

        upper_a_line.SetLineWidth(3);

        upper_a_line.Draw();

 //   }

 //   if(cut_b_lower != TMath::Infinity()){

        // make a lower-boundary line

        TLine lower_b_line(x_min, cut_b_lower, x_max, cut_b_lower);

        lower_b_line.SetLineColor(kBlack);

        lower_b_line.SetLineWidth(3);

        lower_b_line.Draw();

 //   }

 //   if(cut_b_upper != TMath::Infinity()){

        // make a lower-boundary line

        TLine upper_b_line(x_min, cut_b_upper, x_max, cut_b_upper);

        upper_b_line.SetLineColor(kBlack);

        upper_b_line.SetLineWidth(3);

        upper_b_line.Draw();

 //   }

    c.cd(2);

    gPad->Update();

    x_min = gPad->GetUxmin();

    x_max = gPad->GetUxmax();

    y_min = gPad->GetUymin();

    y_max = gPad->GetUymax();

    // Now sort out any cut-indicators:

 //   if(cut_a_lower != TMath::Infinity()){

        // make a lower-boundary line

        TLine mc_lower_a_line(cut_a_lower, y_min, cut_a_lower, y_max);

        mc_lower_a_line.SetLineColor(kBlack);

        mc_lower_a_line.SetLineWidth(3);

        mc_lower_a_line.Draw();

 //   }

 //   if(cut_a_upper != TMath::Infinity()){

        // make a lower-boundary line

        TLine mc_upper_a_line(cut_a_upper, y_min, cut_a_upper, y_max);

        mc_upper_a_line.SetLineColor(kBlack);

        mc_upper_a_line.SetLineWidth(3);

        mc_upper_a_line.Draw();

//    }

 //   if(cut_b_lower != TMath::Infinity()){

        // make a lower-boundary line

        TLine mc_lower_b_line(x_min, cut_b_lower, x_max, cut_b_lower);

        mc_lower_b_line.SetLineColor(kBlack);

        mc_lower_b_line.SetLineWidth(3);

        mc_lower_b_line.Draw();

  //  }

 //   if(cut_b_upper != TMath::Infinity()){

        // make a lower-boundary line

        TLine mc_upper_b_line(x_min, cut_b_upper, x_max, cut_b_upper);

        mc_upper_b_line.SetLineColor(kBlack);

        mc_upper_b_line.SetLineWidth(3);

        mc_upper_b_line.Draw();

 //   }

    c.Print(saveAs.c_str());

}

void ApplyPlotStyle::Hist_TOFvTKU(TH2F* data_hist, TH2F* mc_hist, string saveAs, double mean_dP, double min_dP, double max_dP, int legend_location, bool normalise, string legend1, string legend2){

    /*

     * Stack data and MC hists together onto one plot:

     *   - Normalise to 1

     *   - Set plot style

     *   - Stack

     * Legend (stats box) location:

     *   0: no legend

     *   1: top left

     *   2: top right

     *   3: bottom left

     *   4: bottom right

     */

    double norm_data = data_hist->GetMaximum();

    std::cout << "Max in data hist = " << norm_data << "\n";

    double norm_mc = mc_hist->GetMaximum();

    std::cout << "Max in mc hist = " << norm_mc << "\n";

    if(normalise == true){

   //     if(norm_data > norm_mc){

            data_hist->SetMinimum(0.0);

            data_hist->SetMaximum(norm_data);

            mc_hist->SetMinimum(0.0);

            mc_hist->SetMaximum(norm_data);

     //   }

     /*   else{

            data_hist->SetMinimum(0.0);

            data_hist->SetMaximum(norm_mc);

            mc_hist->SetMinimum(0.0);

            mc_hist->SetMaximum(norm_mc);

        }

      */

    }

    data_hist->GetYaxis()->SetTitleOffset(1.5);

    mc_hist->GetYaxis()->SetTitleOffset(1.5);

    data_hist->SetTitle(legend1.c_str());

    mc_hist->SetTitle(legend2.c_str());

    data_hist->GetXaxis()->SetTitleSize(.05);

    data_hist->GetYaxis()->SetTitleSize(.05);

    data_hist->GetXaxis()->SetLabelSize(.05);

    data_hist->GetYaxis()->SetLabelSize(.05);

    data_hist->GetZaxis()->SetLabelSize(.05);

    mc_hist->GetXaxis()->SetTitleSize(.05);

    mc_hist->GetYaxis()->SetTitleSize(.05);

    mc_hist->GetXaxis()->SetLabelSize(.05);

    mc_hist->GetYaxis()->SetLabelSize(.05);

    mc_hist->GetZaxis()->SetLabelSize(.05);

    TCanvas c("c", "c", 800, 400);

    c.Divide(2, 1);

    c.SetFillColor(0);

    c.SetFrameFillStyle(0);

    c.cd(1);

    data_hist->Draw("colz");

    c.cd(2);

    mc_hist->Draw("colz");

    if(legend_location == 1){

        gStyle->SetStatX(0.33);

        gStyle->SetStatY(0.88);

        gStyle->SetOptStat("emr");

    }

    else if(legend_location == 2){

        gStyle->SetStatX(0.88);

        gStyle->SetStatY(0.88);

        gStyle->SetOptStat("emr");

    }

    else if(legend_location == 3){

        gStyle->SetStatX(0.33);

        gStyle->SetStatY(0.38);

        gStyle->SetOptStat("emr");

    }

    else if(legend_location == 4){

        gStyle->SetStatX(0.88);

        gStyle->SetStatY(0.38);

        gStyle->SetOptStat("emr");

    }

    else{

        gStyle->SetOptStat(0);

    }

    c.cd(1);

    TF1 *f1 = new TF1("testfunc", tofFunc, -1.0, 351.0, 2);

    f1->SetParameter(0, 105.668);

    f1->SetParameter(1, mean_dP);

    TF1 *f2 = new TF1("testfunc", tofFunc, -1.0, 351.0, 2);

    f2->SetParameter(0, 105.668);

    f2->SetParameter(1, min_dP);

    TF1 *f3 = new TF1("testfunc", tofFunc, -1.0, 351.0, 2);

    f3->SetParameter(0, 105.668);

    f3->SetParameter(1, max_dP);

    f1->SetLineStyle(3);

    f1->SetLineColor(kWhite);

    f1->Draw("same,APL");

    f2->SetLineColor(kBlack);

    f2->SetLineWidth(3);

    f2->Draw("same,APL");

    f3->SetLineColor(kBlack);

    f3->SetLineWidth(3);

    f3->Draw("same,APL");

    gPad->Update();

    /*

     TPaletteAxis *palette = (TPaletteAxis*)data_hist->GetListOfFunctions()->FindObject("palette");

     // the following lines moe the paletter. Choose the values you need for the position.

     palette->SetX1NDC(0.9);

     palette->SetX2NDC(0.92);

     palette->SetY1NDC(0.1);

     palette->SetY2NDC(0.9);

     gPad->Modified();

     gPad->Update();

     */

    c.cd(2);

    f1->Draw("same,APL");

    f2->Draw("same,APL");

    f3->Draw("same,APL");

    gPad->Update();

    c.Print(saveAs.c_str());

}

Double_t ApplyPlotStyle::tofFunc(Double_t *x, Double_t *par){

    Double_t m = par[0];

    Double_t dP = par[1];

    double tof1_z = 12929.4396e-3; // m

    double tof0_z = 5285.6636e-3; // m

    double L = tof1_z - tof0_z;

    Double_t p_tku = x[0];

    double p = p_tku + dP;

    double t = L * TMath::Sqrt(m*m + p*p) / (TMath::C()*p);

    return t/1.0e-9;

}

void ApplyPlotStyle::Hist2D_single(TH2F* data_hist, string saveAs="output.pdf", double cut_a_lower=TMath::Infinity(), double cut_a_upper=TMath::Infinity(), double cut_b_lower=TMath::Infinity(), double cut_b_upper=TMath::Infinity(), int legend_location=0,bool normalise=true, string legend1="Data"){

    /*

     * Stack data and MC hists together onto one plot:

     *   - Normalise to 1

     *   - Set plot style

     *   - Stack

     * Legend (stats box) location:

     *   0: no legend

     *   1: top left

     *   2: top right

     *   3: bottom left

     *   4: bottom right

     */

 //   data_hist->GetYaxis()->SetTitleOffset(1.5);

    data_hist->GetXaxis()->SetTitleSize(.05);

    data_hist->GetYaxis()->SetTitleSize(.05);

    data_hist->GetXaxis()->SetLabelSize(.05);

    data_hist->GetYaxis()->SetLabelSize(.05);

    data_hist->GetZaxis()->SetLabelSize(.05);

    data_hist->SetTitle(legend1.c_str());

    TCanvas c("c", "c", 800, 800);

    data_hist->Draw("colz");

    if(legend_location == 1){

        gStyle->SetStatX(0.33);

        gStyle->SetStatY(0.88);

        gStyle->SetOptStat("emr");

    }

    else if(legend_location == 2){

        gStyle->SetStatX(0.88);

        gStyle->SetStatY(0.88);

        gStyle->SetOptStat("emr");

    }

    else if(legend_location == 3){

        gStyle->SetStatX(0.33);

        gStyle->SetStatY(0.38);

        gStyle->SetOptStat("emr");

    }

    else if(legend_location == 4){

        gStyle->SetStatX(0.88);

        gStyle->SetStatY(0.38);

        gStyle->SetOptStat("emr");

    }

    else{

        gStyle->SetOptStat(0);

    }

    gPad->Update();

    double x_min = gPad->GetUxmin();

    double x_max = gPad->GetUxmax();

    double y_min = gPad->GetUymin();

    double y_max = gPad->GetUymax();

    // Now sort out any cut-indicators:

    if(cut_a_lower != TMath::Infinity()){

        // make a lower-boundary line

        TLine lower_a_line(cut_a_lower, y_min, cut_a_lower, y_max);

        lower_a_line.SetLineColor(kBlack);

        lower_a_line.SetLineWidth(3);

        lower_a_line.Draw();

    }

    if(cut_a_upper != TMath::Infinity()){

        // make a lower-boundary line

        TLine upper_a_line(cut_a_upper, y_min, cut_a_upper, y_max);

        upper_a_line.SetLineColor(kBlack);

        upper_a_line.SetLineWidth(3);

        upper_a_line.Draw();

    }

    if(cut_b_lower != TMath::Infinity()){

        // make a lower-boundary line

        TLine lower_b_line(x_min, cut_b_lower, x_max, cut_b_lower);

        lower_b_line.SetLineColor(kBlack);

        lower_b_line.SetLineWidth(3);

        lower_b_line.Draw();

    }

    if(cut_b_upper != TMath::Infinity()){

        // make a lower-boundary line

        TLine upper_b_line(x_min, cut_b_upper, x_max, cut_b_upper);

        upper_b_line.SetLineColor(kBlack);

        upper_b_line.SetLineWidth(3);

        upper_b_line.Draw();

    }

    c.Print(saveAs.c_str());

}