#include "Plot.h"

// ROOT includes
#include "TArrow.h"
#include "TCanvas.h"
#include "TFile.h"
#include "TFrame.h"
#include "TGraphAsymmErrors.h"
#include "TH1.h"
#include "TH1D.h"
#include "THStack.h"
#include "TLatex.h"
#include "TLegend.h"
#include "TMath.h"
#include "TPad.h"
#include "TStyle.h"

// c++ includes
#include <algorithm>
#include <iostream>

using namespace std;

//_____________________________________________________________________________
//
Plot::Plot(std::string name,int canvasWidth,int canvasHeight,bool hideRatioPad){
    fName = name;
    c = new TCanvas(fName.c_str(),fName.c_str(),canvasWidth,canvasHeight);
    //
    if(hideRatioPad){
        pad0 = new TPad("pad0","pad0",0,0,1,1,0,0,0);
    }
    else{
        pad0 = new TPad("pad0","pad0",0,0.20,1,1,0,0,0);
    }
    pad0->SetTicks(1,1);
    pad0->SetTopMargin(0.05*(700./canvasHeight));
    if(hideRatioPad){
        pad0->SetBottomMargin(0.14*(600./canvasHeight));
    }
    else{
        pad0->SetBottomMargin(0.1);
    }
    pad0->SetLeftMargin(0.14*(600./canvasWidth));
    pad0->SetRightMargin(0.05*(600./canvasWidth));
    pad0->SetFrameBorderMode(0);
    pad0->SetFillStyle(0);
    //
    if(hideRatioPad){
        pad1 = nullptr;
    }
    else{
        pad1 = new TPad("pad1","pad1",0,0,1,0.28,0,0,0);
        pad1->SetTicks(1,1);
        pad1->SetTopMargin(0.0);
        pad1->SetBottomMargin(0.37*(700./canvasHeight));
        pad1->SetLeftMargin(0.14*(600./canvasWidth));
        pad1->SetRightMargin(0.05*(600./canvasWidth));
        pad1->SetFrameBorderMode(0);
        pad1->SetFillStyle(0);
    }
    //
    if(pad1!=nullptr) pad1->Draw();
    pad0->Draw();
    pad0->cd();
    h_stack = new THStack("h_stack","h_stack");
    h_tot = nullptr;
    g_tot = nullptr;
    xtitle = "Variable [GeV]";
    ytitle = "Events";
    fLabels.clear();
    fLumi = "20.3 fb^{-1}";
    fCME = "8 TeV";
    fATLASlabel = "none";
    yMaxScale = 2.;
    Chi2val = -1;
    NDF = -1;
    Chi2prob = -1;
    KSprob = -1;
    //
    h_data = nullptr;
    g_data = nullptr;

    h_bkg.clear();
    h_signal.clear();
    h_normsig.clear();
    h_oversig.clear();

    //
    fIsNjet = false;
    fShowYields = false;
    //
    fDataName = "Data";
    fSigNames.clear();
    fNormSigNames.clear();
    fOverSigNames.clear();
    fBkgNames.clear();

    fBinWidth = -1;
    fLumiScale = 1.;
    fBlindingThreshold = -1; // if <0, no blinding
    fLegendNColumns = 2;

    fYmin = 0;
    fYmax = 0;
    fRatioYmin = 0.;
    fRatioYmax = 2.;

    h_blinding = nullptr;

    h_tot_bkg_prefit = nullptr;

    leg = nullptr;
    leg1 = nullptr;

    fRatioYtitle = "";
    fRatioType = "DATA/MC";
    fLabelX = -1;
    fLabelY = -1;
    fLegendX1 = -1;
    fLegendX2 = -1;
    fLegendY = -1;
}

//_____________________________________________________________________________
//
Plot::~Plot(){
    delete h_data;
    h_bkg    .clear();
    h_signal .clear();
    h_normsig.clear();
    h_oversig.clear();
    delete h_stack;
    delete h_tot;
    delete g_tot;
    delete h_blinding;
    delete h_tot_bkg_prefit;
    delete h_dummy;
    if(leg!=nullptr) delete leg;
    if(leg1!=nullptr) delete leg1;
    delete pad0;
    delete pad1;
    delete c;
}

//_____________________________________________________________________________
//
void Plot::SetChannel(const std::string& name){
    fLabels.clear();
    fLabels.push_back(name);
}

//_____________________________________________________________________________
//
void Plot::AddLabel(const std::string& name){
    fLabels.push_back(name);
}

//_____________________________________________________________________________
//
void Plot::SetLumi(const std::string& name){
    fLumi = name;
}

//_____________________________________________________________________________
//
void Plot::SetLumiScale(float scale){
    fLumiScale = scale;
}

//_____________________________________________________________________________
//
void Plot::SetCME(const std::string& name){
    fCME = name;
}

//_____________________________________________________________________________
//
void Plot::SetXaxis(const std::string& name,bool isNjet){
    xtitle = name;
    fIsNjet = isNjet;
}

//_____________________________________________________________________________
//
void Plot::SetYaxis(const std::string& name){
    ytitle = name;
}

//_____________________________________________________________________________
//
void Plot::SetYmaxScale(float scale){
    yMaxScale = scale;
}

//_____________________________________________________________________________
//
void Plot::SetBinLabel(int bin, const std::string& name){
    fBinLabel[bin] = name;
}

//_____________________________________________________________________________
//
void Plot::SetBinWidth(float width){
    fBinWidth = width;
}

//_____________________________________________________________________________
//
void Plot::SetData(TH1* h,std::string name){
    h_data = (TH1*)h->Clone();
    // if no name is given, take the histogram title
    if(name=="") name = h->GetTitle();
    fDataName = name;
}

//_____________________________________________________________________________
//
void Plot::AddSignal(TH1* h,std::string name){
    // if no name is given, take the histogram title
    if(name=="") name = h->GetTitle();
    unsigned int idx = std::find(fSigNames.begin(),fSigNames.end(),name) - fSigNames.begin();
    if(idx<fSigNames.size()){
        h_signal[idx]->Add(h,fLumiScale);
    }
    else{
        h_signal.push_back((TH1*)h->Clone());
        h_signal[fSigNames.size()]->Scale(fLumiScale);
        fSigNames.push_back(name);
    }
}

//_____________________________________________________________________________
//
void Plot::AddBackground(TH1* h,std::string name){
    if(h_tot==nullptr) h_tot = (TH1*)h->Clone();
    else h_tot->Add(h);
    // if no name is given, take the histogram title
    if(name=="") name = h->GetTitle();
    //
    unsigned int idx = std::find(fBkgNames.begin(),fBkgNames.end(),name) - fBkgNames.begin();
    if(idx<fBkgNames.size()){
        h_bkg[idx]->Add(h,fLumiScale);
    }
    else{
        h_bkg.push_back((TH1*)h->Clone());
        h_bkg[fBkgNames.size()]->Scale(fLumiScale);
        fBkgNames.push_back(name);
    }
}