//Script used to make pretty plots of the Q2 veto
//Renata Kopecna

#include "../GlobalFunctions.hh"
#include "../Paths.hpp"
#include "../Utils.hpp"
#include "../Design.hpp"

double phi_low = 0.98;
double phi_up  = 1.1;
double jpsi_low= 8.0;
double jpsi_up = 11.0;
double psi_low = 12.5;
double psi_up  = 15.0;

int makePrettyQ2plot(){

    LHCbStyle();

    //Load the files
    TChain* theTree = new TChain("DecayTree");
    for (auto &yr: yearsData(12)){
        theTree->Add(GetBDTinputFile(yr, false, false, false,false).c_str());
        coutDebug("Reading " + GetBDTinputFile(yr, false, false, false,false));
    }

    //init the histogram
    string hist_name = "hist";
    TH2D* hist = new TH2D(hist_name.c_str(),hist_name.c_str(),200,cut_B_plus_M_low,cut_B_plus_M_high,200,0.1,19.0);

    //Define the plotting options
    string B_branch = UseDTF ? "B_plus_M_DTF" : "B_plus_M";
    string plot_branch = "Q2/1000000:" + B_branch;
    string plot_opts = "COLZ";

    string drawCommand = plot_branch + ">>" + hist_name;
    coutDebug("Drawing this: " + drawCommand);

    //Draw into the histogram
    theTree->Draw(drawCommand.c_str());

    //Design histograms
    design_TH2F(hist,"K^{+}#pi^{0}#mu^{+}#mu^{-} [MeV]","q^{2} [MeV^{2}]","Entries [a.u.]");
    hist->GetXaxis()->SetLabelSize(0.05);
    hist->GetYaxis()->SetLabelSize(0.05);
    hist->GetYaxis()->SetTitleOffset(1.0);
    hist->GetZaxis()->SetTitleOffset(1.2);

    //Add the resonance lines
    TLine *linePhiOne  = design_veto_line(cut_B_plus_M_low,phi_low,cut_B_plus_M_high,phi_low,kRed);
    TLine *linePhiTwo  = design_veto_line(cut_B_plus_M_low,phi_up,cut_B_plus_M_high,phi_up,kRed);
    TLine *lineJspiOne = design_veto_line(cut_B_plus_M_low,jpsi_low,cut_B_plus_M_high,jpsi_low, kRed);
    TLine *lineJspiTwo = design_veto_line(cut_B_plus_M_low,jpsi_up,cut_B_plus_M_high,jpsi_up,kRed);
    TLine *linePsiOne  = design_veto_line(cut_B_plus_M_low,psi_low,cut_B_plus_M_high,psi_low,kRed);
    TLine *linePsiTwo  = design_veto_line(cut_B_plus_M_low,psi_up,cut_B_plus_M_high,psi_up,kRed);

    //Save the histogram
    TCanvas *canvas = c_TH2F("q2_veto");
    canvas->SetRightMargin(0.18);
    canvas->SetLeftMargin(0.12);
    canvas->cd();
    canvas->SetLogz();
    hist->Draw(plot_opts.c_str());
    linePhiOne->Draw("SAME");
    linePhiTwo->Draw("SAME");
    lineJspiOne->Draw("SAME");
    lineJspiTwo->Draw("SAME");
    linePsiOne->Draw("SAME");
    linePsiTwo->Draw("SAME");


    canvas->SaveAs("./q2_veto.eps");
    canvas->SaveAs("./q2_veto.root");

    return 1;

}

int makePrettyQ2dist(){

    LHCbStyle();

    //Load the files
    TChain* theTree = new TChain("DecayTree");
    for (auto &yr: yearsData(12)){
        theTree->Add(GetBDTinputFile(yr, false, false, false,false).c_str());
        coutDebug("Reading " + GetBDTinputFile(yr, false, false, false,false));
    }

    //init the histogram
    string hist_name = "hist";
    TH1D* hist = new TH1D(hist_name.c_str(),hist_name.c_str(),200,0.1,19.0);

    //Define the plotting options
    string B_branch = UseDTF ? "B_plus_M_DTF" : "B_plus_M";
    string plot_branch = "Q2/1000000";
    string plot_opts = "";
    string plot_cut = to_string(cut_B_plus_M_low) +"<"+ B_branch +" &&  "+ B_branch +"<"+ to_string(cut_B_plus_M_high);

    string drawCommand = plot_branch + ">>" + hist_name;
    coutDebug("Drawing this: " + drawCommand);
    coutDebug("with a cut: " + plot_cut);

    //Draw into the histogram
    theTree->Draw(drawCommand.c_str(),plot_cut.c_str());

    //Design histograms
    design_lines(hist,"q^{2} [GeV^{2}]","Entries (a.u.)", kAzure+2);
    hist->GetXaxis()->SetLabelSize(0.05);
    hist->GetYaxis()->SetLabelSize(0.05);
    hist->GetYaxis()->SetRangeUser(20,260000);
    hist->GetYaxis()->SetTitleOffset(1.2);

    //Add the resonance lines
    TH1D* histPhi  = new TH1D("phi", "phi", 1, phi_low, phi_up);
    TH1D* histJpsi = new TH1D("jpsi","jspi",1, jpsi_low,jpsi_up);
    TH1D* histPsi  = new TH1D("psi", "psi", 1, psi_low, psi_up);

    histPhi->SetBinContent(1,260000);
    histJpsi->SetBinContent(1,260000);
    histPsi->SetBinContent(1,260000);

    //TLine *linePhiOne  = design_veto_line(phi_low, 20, phi_low, 260000, kRed);
    //TLine *linePhiTwo  = design_veto_line(phi_up, 20, phi_up, 260000, kRed);
    //TLine *lineJspiOne = design_veto_line(jpsi_low , 20, jpsi_low,  260000, kRed);
    //TLine *lineJspiTwo = design_veto_line(jpsi_up, 20, jpsi_up, 260000, kRed);
    //TLine *linePsiOne  = design_veto_line(psi_low, 20, psi_low, 260000, kRed);
    //TLine *linePsiTwo  = design_veto_line(psi_up, 20, psi_up, 260000, kRed);

    string main_name = "q2_dist";
    histPhi->SetFillStyle(3444); //3344
    histPhi->SetFillColor(kRed+1);
    histPhi->SetLineWidth(0);
    histJpsi->SetFillStyle(3444); //3344
    histJpsi->SetFillColor(kRed+1);
    histJpsi->SetLineWidth(0);
    histPsi->SetFillStyle(3444); //3344
    histPsi->SetFillColor(kRed+1);
    histPsi->SetLineWidth(0);

    //Save the histogram
    TCanvas *canvas = c_canvas(main_name);
    canvas->SetLeftMargin(0.12);
    canvas->cd();
    canvas->SetLogy();
    hist->Draw(plot_opts.c_str());
    histPhi->Draw("SAME");
    histJpsi->Draw("SAME");
    histPsi->Draw("SAME");

    //linePhiOne->Draw("SAME");
    //linePhiTwo->Draw("SAME");
    //lineJspiOne->Draw("SAME");
    //lineJspiTwo->Draw("SAME");
    //linePsiOne->Draw("SAME");
    //linePsiTwo->Draw("SAME");

    canvas->SaveAs(Form("%s.eps",main_name.c_str()));
    canvas->SaveAs(Form("%s.root",main_name.c_str()));

    return 1;

}