#include <TTree.h>
#include <TPad.h>
#include <TLegend.h>
#include <TLegendEntry.h>
#include <TROOT.h>
#include <TFile.h>
#include <TCanvas.h>
#include <TH2F.h>
#include <TStyle.h>
#include <TChain.h>
#include <TColor.h>
#include <TPolyLine.h>
#include <TCutG.h>

#include <fstream>
#include <iostream>
#include <string>

using namespace std;

// INSTRUCTIONS:
//
// This macro generates a 2D histogram of the transverse momentum vs Z position
// of particles hitting the CSI calorimeter. It also prints out at the end where 
//
// To run this macro use the following form: ./PTvsRecZplot.exe [ARG_1] [ARG_2]
//
// Where ARG_1 is the path to the root file/files you wish to evaluate and ARG_2
// is the path to save the generated histogram.
//
// If the data you wish to evaluate is spread between several root files, use
// regular expression to group them within a single argument in ARG_1.

int main(int argc, char** argv) {
  // Created canvas for histogram
  TCanvas *c1 = new TCanvas("c1", "c1", 900, 500);

  // Initialize chain to hold tree variables
  TChain *chain = new TChain("RecTree");

  // Initialized histogram
  TH2F *Pi0PtvsPi0RecZ = new TH2F("Pi0PtvsPi0RecZ",
								  "#pi^{0} P_{t} vs #pi^{0} Z_{vtx}",
								  100, 1500, 6000, 100, 0, 500);

  // Array used for initializing signal box and signal perimeter
  Double_t xCoordSignal[] = {3000, 4000, 4700, 4700, 3000, 3000};
  Double_t yCoordSignal[] = {130, 130, 150, 250, 250, 130};

  Double_t xCoordBlind[] = {2900, 5100, 5100, 2900, 2900};
  Double_t yCoordBlind[] = {120, 120, 260, 260, 120};

  Double_t xCoordLower[] = {2900, 5100, 5100, 2900, 2900};
  Double_t yCoordLower[] = {0, 0, 120, 120, 0};

  // Initialized Boxes
  TPolyLine* signalRegion = new TPolyLine(6, xCoordSignal, yCoordSignal);
  TPolyLine* blindedRegion = new TPolyLine(5, xCoordBlind, yCoordBlind);
  TPolyLine* lowerPtRegion = new TPolyLine(5, xCoordLower, yCoordLower);

  // Used to mark borders for counting the number of events in a region
  TCutG* signalPerimeter = new TCutG("", 6, xCoordSignal, yCoordSignal);
  TCutG* blindedPerimeter = new TCutG("", 5, xCoordBlind, yCoordBlind);
  TCutG* lowerPtPerimeter = new TCutG("", 5, xCoordLower, yCoordLower);

 // Number of events in each region
  Double_t countSignalRegion = 0;
  Double_t countBlindedRegion = 0;
  Double_t countLowerPtRegion = 0;

  // Easier to manage several root files at once
  chain->Add(argv[1]);

  // Set style of boxes
  signalRegion->SetFillColorAlpha(0, 0.);
  signalRegion->SetFillStyle(0);
  signalRegion->SetLineWidth(2);
  signalRegion->SetLineColor(1);
  
  blindedRegion->SetFillColorAlpha(0, 0.);
  blindedRegion->SetFillStyle(0);
  blindedRegion->SetLineWidth(1);
  blindedRegion->SetLineColor(1);
  
  lowerPtRegion->SetFillColorAlpha(0, 0.);
  lowerPtRegion->SetFillStyle(0);
  lowerPtRegion->SetLineWidth(1);
  lowerPtRegion->SetLineColor(1);

  // Draw data of chain into histogram with no cuts
  chain->Draw("Pi0Pt:Pi0RecZ>>Pi0PtvsPi0RecZ");

  // Apply style to histogram and draw it
  Pi0PtvsPi0RecZ->Draw("COLZ");

  // Add X axis label and center it
  Pi0PtvsPi0RecZ->GetXaxis()->SetTitle("#pi^{0} Z_{vtx} [mm]");
  Pi0PtvsPi0RecZ->GetXaxis()->CenterTitle(true);

  // Add Y axis label and center it
  Pi0PtvsPi0RecZ->GetYaxis()->SetTitle("#pi^{0} P_{t} [MeV/c]");
  Pi0PtvsPi0RecZ->GetYaxis()->CenterTitle(true);

  // Compute count in signalRegion
  countSignalRegion = signalPerimeter->IntegralHist(Pi0PtvsPi0RecZ);

  // Compute count in blindedRegion
  countBlindedRegion = blindedPerimeter->IntegralHist(Pi0PtvsPi0RecZ);

  // Compute count in lowerPtRegion
  countLowerPtRegion = lowerPtPerimeter->IntegralHist(Pi0PtvsPi0RecZ);

  // Draw boxes
  signalRegion->Draw();
  blindedRegion->Draw();
  lowerPtRegion->Draw();

  // Stat Box options specified
  gStyle->SetOptStat(0);

  cout << "#Events in Signal Region: " << countSignalRegion << endl;
  cout << "#Events in Blinded Region: " << countBlindedRegion << endl;
  cout << "#Events in Lower Pt Region: " << countLowerPtRegion << endl;

  // Save histogram as ARG_2
  c1->Print(argv[2]);

  delete chain;
  delete Pi0PtvsPi0RecZ;
  delete signalRegion;
  delete blindedRegion;
  delete lowerPtRegion;
  delete c1;

  return 0;
}