//MoEDAL Selector

#include "TROOT.h"
#include "TString.h"

#include "TChain.h"
#include "TFile.h"
#include "TCanvas.h"
#include "TSelector.h"
#include "TPad.h"
#include "TH2F.h"
#include "TH3F.h"
#include "TTreeReader.h"
#include "TTreeReaderArray.h"
#include "TStyle.h"
#include "TattAxis.h"

class MoEDALSelectorHist : public TSelector {
public :
    
   //Define all varibales require to run this code
    
    const char *outputfile;
 //  TTree *outputtree = new TTree("SleptonNtuple", "Slepton Simulation Data");
   //As in Histogramming, we want to histogram the particles' mc_initial_px; the type of the data member should thus be TH1F*. 
   Int_t j;
   Int_t j1;
   Int_t j2;
   Int_t j3;
   Int_t j4;
   Int_t j5;
   Float_t weight;
   
   TH2F *hist_px_py;
   TH2F *hist_x_y;
   TH3F *hist_x_y_z;
   TH2F *hist_x_y_mmt;
   TH2F *hist_x_y_ntd;
   TH2F *hist_x_y_ntd2;
   TH2F *hist_x_y_ntddisp;
    TH3F *hist_ntd_x_y_z;
   //TH2F *hist_px_py_new = (TH2F*)hist_px_py->Clone(); //cloning a histogram
   
   TTreeReader fReader;
   TTreeReaderArray<Float_t> mc_initial_px;
   TTreeReaderArray<Float_t> mc_initial_py;
   TTreeReaderArray<Float_t> mc_final_x;
   TTreeReaderArray<Float_t> mc_final_y;
   TTreeReaderArray<Float_t> mc_final_z;
   TTreeReaderArray<Float_t> mmthit_hit_x;
   TTreeReaderArray<Float_t> mmthit_hit_y;
   TTreeReaderArray<Float_t> ntdhit_hit_entryx;
   TTreeReaderArray<Float_t> ntdhit_hit_entryy;
   TTreeReaderArray<Float_t> ntdhit_hit_entryz;
   TTreeReaderArray<Float_t> ntdhit_hit_dispx;
   TTreeReaderArray<Float_t> ntdhit_hit_dispy;
   TTreeReaderArray<Float_t> ntdhit_hit_dispz;

   //Don't forget to initialize the histogram pointer to 0 in the constructor.
   MoEDALSelectorHist(TTree * = 0): j(0), j1(0), j2(0), j3(0), j4(0), j5(0), weight(1), hist_px_py(0), hist_x_y(0), hist_x_y_mmt(0), hist_x_y_ntd(0),hist_x_y_ntd2(0), hist_x_y_ntddisp(0), hist_x_y_z(0),
                      mc_initial_px(fReader, "mc_initial_px"), mc_initial_py(fReader, "mc_initial_py"), mc_final_x(fReader, "mc_final_x"), mc_final_y(fReader, "mc_final_y"), mc_final_z(fReader, "mc_final_z"), mmthit_hit_x(fReader, "mmthit_hit_x"), mmthit_hit_y(fReader, "mmthit_hit_y"),
                      ntdhit_hit_entryx(fReader, "ntdhit_hit_entryx"), ntdhit_hit_entryy(fReader, "ntdhit_hit_entryy"), ntdhit_hit_entryz(fReader, "ntdhit_hit_entryz"),ntdhit_hit_dispx(fReader, "ntdhit_hit_dispx"), ntdhit_hit_dispy(fReader, "ntdhit_hit_dispy"), ntdhit_hit_dispz(fReader, "ntdhit_hit_dispz"){ }
   virtual ~MoEDALSelectorHist() { }

   virtual void    Init(TTree *tree);
   virtual void    Begin(TTree *tree);
   virtual void    SlaveBegin(TTree *tree);
   virtual Bool_t  Process(Long64_t entry);
   virtual void    Terminate();
    // Determines how the TSelector operates
   virtual Int_t   Version() const { return 2; }
    
    ClassDef(MoEDALSelectorHist,0);
};

// Set branch adresses
void MoEDALSelectorHist::Init(TTree *tree)
{
    
    fReader.SetTree(tree);
}

// call all local graphics and file systems
void MoEDALSelectorHist::Begin() {
    
}

// All the initialization needed for Proccess() happens here
void MoEDALSelectorHist::SlaveBegin(TTree *tree)
{
    
   //initialize histograms
   hist_px_py = new TH2F("hist_px_py","Phase space distribution of monopoles generated",50,-600000,600000,50,-600000,600000);
   hist_x_y = new TH2F("hist_x_y","Spacial distribution of final state monopoles",50,-3000,3000,50,-3000,3000);
   hist_x_y_z = new TH3F("hist_x_y_z","Spacial distribution of final state monopoles",50,-3000,3000,50,-6000,-4900,50,-3000,3000);
   hist_x_y_mmt = new TH2F("hist_x_y_mmt","Spacial distribution of monopoles hits in the MMT",50,-3000,3000,50,-3000,3000);
   hist_x_y_ntd = new TH2F("hist_x_y_ntd","Spacial distribution of monopoles entering the NTD",50,-3000,3000,50,-3000,3000);
   hist_x_y_ntd2 = new TH2F("hist_x_y_ntd","Spacial distribution of monopoles entering the NTD",50,-3000,3000,50,-3000,3000);
   hist_x_y_ntddisp = new TH2F("hist_x_y_ntddisp","Spacial displacement of monopoles in the NTD",50,-0.5,0.5,50,-0.5,0.5);
   hist_ntd_x_y_z = new TH3F("hist_ntd_x_y_z","Spacial distribution of slepton hits in NTDs",50,-3000,1500,50,-2500,0, 50,-2000,1500);
   //hist_px_py_new = (TH2F*)hist_px_py->Clone(); //cloning a histogram
    //50,-1500,600,50,-1500,0);
   //activate errors
   hist_px_py->Sumw2();

   //TProfile *hist_px_py = new TProfile("hist_px_py","Hist_px_py",1000000,-200000,200000,-200000,200000);

}

Bool_t MoEDALSelectorHist::Process(Long64_t entry)
{
    
    fReader.SetEntry(entry);
    
    //Get the data from the current TTree entry by getting
    //the value from the connected reader
    
    //loop over each event vector, there are 10 events of varying sizes
    for (int i=0; i<mc_initial_px.GetSize(); i++ ){
        hist_px_py->Fill(mc_initial_px[i], mc_initial_py[i], weight);
        j1++;
    }
    for (int k=0; k<mc_final_x.GetSize(); k++ ){
        hist_x_y->Fill(mc_final_x[k], mc_final_y[k], weight);
        //printf("first position values: %f   %f  \n", mc_final_x[k], mc_final_y[k]);
        j2++;
    }
    for (int k=0; k<mc_final_z.GetSize(); k++ ){
        hist_x_y_z->Fill(mc_final_z[k], mc_final_x[k], mc_final_y[k]);
        if (mc_final_z[k]>= 5000) printf("second position values: %f   %f   %f \n", mc_final_x[k], mc_final_y[k], mc_final_z[k]);
    }
    for (int l=0; l<mmthit_hit_x.GetSize(); l++ ){
        hist_x_y_mmt->Fill(mmthit_hit_x[l], mmthit_hit_y[l], weight);
    }
    for (int s=0; s<ntdhit_hit_entryx.GetSize(); s++ ){
        hist_x_y_ntd->Fill(ntdhit_hit_entryx[s], ntdhit_hit_entryy[s], weight);
        j3++;
    }
    for (int s=0; s<ntdhit_hit_entryx.GetSize(); s++ ){
        Float_t a = ntdhit_hit_entryx[s]+ntdhit_hit_dispx[s];
        Float_t b = ntdhit_hit_entryy[s]+ntdhit_hit_dispy[s];
        hist_x_y_ntd2->Fill(a, b, weight);
        j4++;
    }
    for (int s=0; s<ntdhit_hit_entryx.GetSize(); s++ ){
        Float_t a = ntdhit_hit_entryx[s]+ntdhit_hit_dispx[s];
        Float_t b = ntdhit_hit_entryy[s]+ntdhit_hit_dispy[s];
        Float_t c = ntdhit_hit_entryz[s]+ntdhit_hit_dispz[s];
        hist_ntd_x_y_z->Fill(c, a, b);
        if (a>= 1500 && a<= -1000) printf("a %f \n", a);
    
        if (b>= 0 && b<= -1000) printf("b %f \n", b);

        if (c>= 0 && c<= -1000) printf("c %f \n", c);

    }
    for (int s=0; s<ntdhit_hit_dispx.GetSize(); s++ ){
        hist_x_y_ntddisp->Fill(ntdhit_hit_dispx[s], ntdhit_hit_dispy[s], weight);
        j5++;
    }
    //Event counter
    j++;

    return kTRUE;
}

void MoEDALSelectorHist::Terminate()
{
    //Styles
    //loose the stats box
    gStyle->SetOptStat(0);
    
    //Check number of Events - j prints 10 if correct
    printf("I made it to %d Events!\n c1: %d \n c2: %d \n c3: %d \n c4: %d \n c5: %d \n", j,j1,j2,j3,j4,j5);
    
    //initialize a Canvas - the canvas was moved here from above the histograms booking stage as the canvas has to be genrated and used in one sweep - see https://root.cern.ch/canvases-and-pads
    TCanvas* canvas_1 = new TCanvas("canvas_1","canvas_1");
    TCanvas* canvas_2 = new TCanvas("canvas_2","canvas_2");
    TCanvas* canvas_3 = new TCanvas("canvas_3","canvas_3");
    TCanvas* canvas_4 = new TCanvas("canvas_4","canvas_4");
    TCanvas* canvas_5 = new TCanvas("canvas_5","canvas_5");
    TCanvas* canvas_6 = new TCanvas("canvas_6","canvas_6");
    TCanvas* canvas_7 = new TCanvas("canvas_7","canvas_7");
    //canvas_1->Divide(1,2);
    
    //Naming the axes
    hist_px_py->GetXaxis()->SetTitle("momentum x (MeV)");
    hist_px_py->GetYaxis()->SetTitle("momentum y (MeV)");
    hist_px_py->GetZaxis()->SetTitle("counts in number of particles");
    ///hist_px_py->GetXaxis()->SetNoExponent(kTRUE);
    hist_px_py->SetLabelSize(0.03,"mc_initial_px");
    hist_px_py->SetLabelSize(0.03,"mc_initial_py");
    hist_px_py->SetLabelSize(0.03,"Z");
    hist_px_py->GetXaxis()->SetLabelOffset();
    hist_px_py->GetYaxis()->SetLabelOffset();
    hist_px_py->GetZaxis()->SetLabelOffset();
    hist_px_py->GetXaxis()->SetTitleOffset(1.5);
    hist_px_py->GetYaxis()->SetTitleOffset(2);
    hist_px_py->GetZaxis()->SetTitleOffset(1);
    hist_x_y->GetXaxis()->SetTitle("position x (mm)");
    hist_x_y->GetYaxis()->SetTitle("position y (mm)");
    hist_x_y->GetZaxis()->SetTitle("counts in number of particles");
    hist_x_y->GetXaxis()->SetTitleOffset(1);
    hist_x_y->GetYaxis()->SetTitleOffset(1.3);
    hist_x_y->GetZaxis()->SetTitleOffset(1);
    
    hist_x_y_z->GetXaxis()->SetTitle("position z (mm)");
    hist_x_y_z->GetYaxis()->SetTitle("position x (mm)");
    hist_x_y_z->GetZaxis()->SetTitle("position y (mm)");
    hist_x_y_z->GetXaxis()->SetTitleOffset(1.7);
    hist_x_y_z->GetYaxis()->SetTitleOffset(2);
    hist_x_y_z->GetZaxis()->SetTitleOffset(1.6);
    
    hist_ntd_x_y_z->GetXaxis()->SetTitle("position z (mm)");
    hist_ntd_x_y_z->GetYaxis()->SetTitle("position x (mm)");
    hist_ntd_x_y_z->GetZaxis()->SetTitle("position y (mm)");
    hist_ntd_x_y_z->GetXaxis()->SetTitleOffset(1.7);
    hist_ntd_x_y_z->GetYaxis()->SetTitleOffset(2);
    hist_ntd_x_y_z->GetZaxis()->SetTitleOffset(1.6);
    
    hist_x_y_mmt->GetXaxis()->SetTitle("position x (mm)");
    hist_x_y_mmt->GetYaxis()->SetTitle("position y (mm)");
    hist_x_y_mmt->GetZaxis()->SetTitle("counts in number of particles");
    hist_x_y_mmt->GetXaxis()->SetTitleOffset(1);
    hist_x_y_mmt->GetYaxis()->SetTitleOffset(1.3);
    hist_x_y_mmt->GetZaxis()->SetTitleOffset(1.5);
    hist_x_y_ntd->GetXaxis()->SetTitle("position x (mm)");
    hist_x_y_ntd->GetYaxis()->SetTitle("position y (mm)");
    hist_x_y_ntd->GetZaxis()->SetTitle("counts in number of particles");
    hist_x_y_ntd->GetXaxis()->SetTitleOffset(1);
    hist_x_y_ntd->GetYaxis()->SetTitleOffset(1.3);
    hist_x_y_ntd->GetZaxis()->SetTitleOffset(1.5);
    hist_x_y_ntddisp->GetXaxis()->SetTitle("position x (mm)");
    hist_x_y_ntddisp->GetYaxis()->SetTitle("position y (mm)");
    hist_x_y_ntddisp->GetZaxis()->SetTitle("counts in number of particles");
    hist_x_y_ntddisp->GetXaxis()->SetTitleOffset(1);
    hist_x_y_ntddisp->GetYaxis()->SetTitleOffset(1);
    hist_x_y_ntddisp->GetZaxis()->SetTitleOffset(1);
    
    // Draw histograms
    
    //loose the stats box
    gStyle->SetOptStat(1);
    canvas_1->cd();
    hist_px_py->Draw("surf2");
    //canvas_1->cd(2);
    //hist_px_py_new->Draw();
    canvas_2->cd();
    //hist_x_y->SetFillColor(4);
    hist_x_y->Draw("surf2");
    canvas_6->cd();
    hist_x_y_z->Draw("iso");
    canvas_3->cd();
    hist_x_y_mmt->SetFillColor(2);
    hist_x_y_mmt->Draw("box");
    canvas_4->cd();
    hist_x_y_ntd->SetFillColor(2);
    //hist_x_y_ntd2->SetFillColor(4);
    hist_x_y_ntd->Draw("box");
    hist_x_y_ntd2->Draw("box same");
    canvas_5->cd();
    hist_x_y_ntddisp->Draw("box");
    canvas_7->cd();
    hist_ntd_x_y_z->Draw();
    
    
 //   outputfile->cd();
 //   canvas_1->Write(outputfile);
 //   outputfile->Close();
    
};