#include <cstdlib>
#include <iostream>
#include <fstream>
#include <sstream>
#include <map>
#include <string>
#include <vector>

#include "TChain.h"
#include "TFile.h"
#include "TTree.h"
#include "TH1F.h"
#include "TCanvas.h"
#include "TColor.h"
#include "limits.h"
#include "TLatex.h"
#include "TString.h"
#include "TObjString.h"
#include "TSystem.h"
#include "TROOT.h"

using namespace std;

void PlotData(TString sigFilename, TString bckFilename, TString NTupleName, TString outputFilename, TString *varsName, const int nvars, TString *varsLabel, TString *units, double *sf, double* XlimitsLower, double* XlimitsUpper) {

// read in the signal and background files and get the Trees out of them
  TFile *sigFile = new TFile(sigFilename);
  TObject *sigTuple_object = sigFile->Get(NTupleName);
  TTree *sigTree  = (TTree*)sigTuple_object;

  TFile *bckFile = new TFile(bckFilename);
  TObject *bckTuple_object = bckFile->Get(NTupleName);
  TTree *bckTree  = (TTree*)bckTuple_object;

// get number of data points
  int nSig = sigTree->GetEntries();
  int nBck = bckTree->GetEntries();
  int nEntries[] = {nSig, nBck};

// create arrays to store variables from the Ntuple
  float vars[nvars][2];
  float limits[nvars][2];
  double Ylimits[nvars];

    cout << "Limits" << endl;
  if (outputFilename.Contains("Test") == true ) {
    double lim[] = {5000, 6000, 10000, 3000, 3000, 3000};
    for ( int i=0 ; i<nvars ; i++) { Ylimits[i] = lim[i] ;}
  }
  else if (outputFilename.Contains("Train") == true ) {
    double lim[] = {5000, 6000, 10000, 3000, 3000, 3000};
    for ( int i=0 ; i<nvars ; i++) { Ylimits[i] = lim[i] ;}
  }
  else {
    double lim[] = {10000, 12000, 18000, 5000, 6000, 6000};
    for ( int i=0 ; i<nvars ; i++) { Ylimits[i] = lim[i] ;}
  }

// set the address of the branches for the NTuples so the varibles can be used
  for ( int i=0 ; i<nvars ; i++){
    sigTree->SetBranchAddress("SelVars_Nominal_" + varsName[i] , &vars[i][0]);
    bckTree->SetBranchAddress("SelVars_Nominal_" + varsName[i] , &vars[i][1]);
  }

// getting the x limits for the graphs
  float sigMin, sigMax, bckMin, bckMax;
  for ( int i=0 ; i<nvars ; i++){
    sigMin = sigTree->GetMinimum("SelVars_Nominal_" + varsName[i]);
    sigMax = sigTree->GetMaximum("SelVars_Nominal_" + varsName[i]);
    bckMin = bckTree->GetMinimum("SelVars_Nominal_" + varsName[i]);
    bckMax = bckTree->GetMaximum("SelVars_Nominal_" + varsName[i]);

    if ( sigMin < bckMin ) { limits[i][0] = sigMin; } else { limits[i][0] = bckMin; }
    if ( sigMax > bckMax ) { limits[i][1] = sigMax; } else { limits[i][1] = bckMax; }

    cout << varsName[i] << " limits: {" << limits[i][0] << ", " << limits[i][1] << "}" << endl;

  }

  double cutVals[nvars][2];
// starting value
  cutVals[0][0] = 63.;		  cutVals[0][1] = 141.;		// mBB
  cutVals[1][0] = 60.;		  cutVals[1][1] = 140.;		// mMMC
  cutVals[2][0] = 450.;		  cutVals[2][1] = 3.93921e+06;	// mHH
  cutVals[3][0] = 0.4;		  cutVals[3][1] = 1.;		// drBB
  cutVals[4][0] = 0.2;		  cutVals[4][1] = 1.5;		// drLepTau
  cutVals[5][0] = 0.;	      cutVals[5][1] = 40.;		// mtLepMet

// create output file for histograms with signal and background on the same graph
  TFile* outputFile = new TFile(outputFilename, "RECREATE");

  TCanvas *c[nvars];
  THStack *hs[nvars];
  TH1F *varHist[nvars][2];
  TH1F *cutHist[nvars][2];
  TGaxis *axis[nvars];
  TLegend *leg[nvars];
  TImage *img[nvars];

  for ( int i=0 ; i<nvars ; i++){
    cout << endl << endl << "Starting " << varsName[i] << endl;

    c[i] = new TCanvas("c" + varsName[i], varsName[i]);
    hs[i] = new THStack("hs" + varsName[i], varsName[i] + ";" + varsLabel[i] + " " + units[i] + ";Events");
    varHist[i][0] = new TH1F("hSig" + varsName[i], varsName[i] + " , signal", 100, XlimitsLower[i], XlimitsUpper[i]);
    varHist[i][1] = new TH1F("hBck" + varsName[i], varsName[i] + " , background", 100, XlimitsLower[i], XlimitsUpper[i]);
    cutHist[i][0] = new TH1F("hSig" + varsName[i] + "_afterCuts", varsName[i] + " , signal after cuts", 100, XlimitsLower[i], XlimitsUpper[i]);
    cutHist[i][1] = new TH1F("hBck" + varsName[i] + "_afterCuts", varsName[i] + " , background after cuts", 100, XlimitsLower[i], XlimitsUpper[i]);

    for ( int k=0 ; k<2 ; k++){
      for ( int j=0 ; j<nEntries[k] ; j++){
        if ( k==0 ) {sigTree->GetEntry(j);}
        if ( k==1 ) {bckTree->GetEntry(j);}
        varHist[i][k]->Fill(vars[i][k]*sf[i]);

        if ( vars[0][k]*sf[0] > cutVals[0][0] && vars[0][k]*sf[0] < cutVals[0][1] && vars[1][k]*sf[1] > cutVals[1][0] && vars[1][k]*sf[1] < cutVals[1][1] && vars[2][k]*sf[2] > cutVals[2][0] && vars[2][k]*sf[2] < cutVals[2][1] && vars[3][k]*sf[3] > cutVals[3][0] && vars[3][k]*sf[3] < cutVals[3][1] && vars[4][k]*sf[4] > cutVals[4][0] && vars[4][k]*sf[4] < cutVals[4][1] && vars[5][k]*sf[5] > cutVals[5][0] && vars[5][k]*sf[5] < cutVals[5][1] ) { 

          cout << vars[i][k]*sf[i] << endl;
          cutHist[i][k]->Fill(vars[i][k]*sf[i]);

        }
      }
    }

// making the histos pretty
    varHist[i][0]->SetLineColor(kRed);
    varHist[i][0]->SetLineStyle(kSolid);
    varHist[i][0]->SetLineWidth(2);

    varHist[i][1]->SetLineColor(kBlack);
    varHist[i][1]->SetLineStyle(kSolid);
    varHist[i][1]->SetFillColor(kYellow);
    varHist[i][1]->SetFillStyle(1001);

    cutHist[i][0]->SetLineColor(kBlue);
    cutHist[i][0]->SetLineStyle(kSolid);
    cutHist[i][0]->SetLineWidth(2);

    cutHist[i][1]->SetLineColor(kBlack);
    cutHist[i][1]->SetLineStyle(kSolid);
    cutHist[i][1]->SetFillColor(kGreen);
    cutHist[i][1]->SetFillStyle(1001);

// save the histograms to the outputfile
    varHist[i][0]->Write();
    varHist[i][1]->Write();
    cutHist[i][0]->Write();
    cutHist[i][1]->Write();

    cout << "Y limit = " << Ylimits[i] << endl;

// add the histograms to the stack
    hs[i]->Add(varHist[i][1]);
    hs[i]->Add(varHist[i][0]);
    hs[i]->Add(cutHist[i][1]);
    hs[i]->Add(cutHist[i][0]);
    //hs[i]->GetXaxis()->SetRangeUser(Xlimits[i][0], Xlimits[i][1]);
    hs[i]->SetMaximum(Ylimits[i]);
    hs[i]->Draw("no stack hist");

// making the legend
    leg[i] = new TLegend(0.675, 0.725, 0.875, 0.875);
    leg[i]->SetBorderSize(1);
    leg[i]->SetFillStyle(0);
    leg[i]->SetTextSize(0.03);
    leg[i]->SetHeader(varsName[i] + ":");
    leg[i]->AddEntry(varHist[i][0], "Signal", "l");
    leg[i]->AddEntry(varHist[i][1], "#it{t#bar{t}} Background", "f");
    leg[i]->AddEntry(cutHist[i][0], "Signal after cuts", "l");
    leg[i]->AddEntry(cutHist[i][1], "#it{t#bar{t}} Background after cuts", "f");
    leg[i]->Draw();

    c[i]->Update();
    c[i]->Write();
  }

// save and close the outputfile
  outputFile->Write();
  outputFile->Close();

  cout<< "Output file has been made."<< endl;

  for ( int i=0 ; i<nvars ; i++){
    img[i] = TImage::Create();
    img[i]->FromPad(c[i]);
    if (outputFilename.Contains("Test") == true ) {img[i]->WriteImage(Form("./Cuts/Multiple Variable Cuts/%i. ", i+1) + varsName[i] + " (Test Events after Cuts).png");}
    else if (outputFilename.Contains("Train") == true ) {img[i]->WriteImage(Form("./Cuts/Multiple Variable Cuts/%i. ", i+1) + varsName[i] + " (Train Events after Cuts).png");}
    else {img[i]->WriteImage(Form("./Cuts/Multiple Variable Cuts/%i. ", i+1) + varsName[i] + " (Events after Cuts).png");}

  }

  delete sigTree;
  delete bckTree;

  for ( int i=0 ; i<nvars ; i++){
    delete img[i];
    delete c[i];
  }
 
}

void SLT_NR_VariablePlots_PostCuts() {

  TString sigFilename = "hhttbb.root";
  TString bckFilename = "ttbar_new.root";
  TString NTupleName = "CxAODTuple_Nominal";
  TString outputFilename = "SLT_NonResonantPlots_AfterCuts.root";

  TString varsName[] = {"mBB", "mMMC","mHH", "drBB", "drLepTau", "mtLepMet"};
  const int nvars = 6;
  
  TString varsLabel[] = {"m_{bb}", "m^{MMC}_{#tau#tau}", "m_{hh}", "#DeltaR(b,b)", "#DeltaR(#tau,#tau)", "m^{W}_{T}"};

  TString units[] = {"[GeV]", "[GeV]", " [GeV]", " ", " ", "[GeV]"};

  double sf[] = {1e-3, 1., 1e-3, 1., 1., 1e-3};
  //double Xlimits[nvars][2] = {{0., 500.}, {0., 500.}, {200., 1200.}, {0., 5}, {0., 5}, {0., 200.}};
  double XlimitsLower[] = {0., 0., 200., 0., 0., 0.};
  double XlimitsUpper[] = {500., 500., 1200., 5., 5., 200.};

  PlotData(sigFilename, bckFilename, NTupleName, outputFilename, varsName, nvars, varsLabel, units, sf, XlimitsLower, XlimitsUpper);
  //PlotData("hhttbb_Test.root" , "ttbar_new_Test.root", NTupleName + "_Test", "SLTResonantPlots_Test.root", varsName, nvars, varsLabel, units, sf, Xlimits);
  //PlotData("hhttbb_Train.root", "ttbar_new_Train.root", NTupleName + "_Train", "SLTResonantPlots_Train.root", varsName, nvars, varsLabel, units, sf, Xlimits);

}