//Plot the features of a ICS photon beam phase-space 

#include "TFile.h"
#include "TTree.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TH1I.h"
#include "TF1.h"
#include "TProfile2D.h"
#include "TGraph.h"
#include "TGraph2D.h"
#include "TGraphErrors.h"
#include "TGaxis.h"
#include "TMath.h"
#include "TRandom3.h"
#include "TCanvas.h"
#include "TLegend.h"
#include <iostream>
#include <fstream>
#include <vector>
#include "string.h"
#include <sstream>

using namespace std;

//----------------------------------------------------------------------------------------------

//rescale function
void rescaleXaxis(TH1D* h, Double_t ScaleFactor)
{
	Double_t Xi = 0;
	Char_t binlabel[128]; 
	Int_t bin_index = 0;
	Int_t i = 0;
		
	Double_t step = 1/ScaleFactor;
  	Double_t nbins = h->GetSize()-2;
  	TAxis* xax = h->GetXaxis();	
  	Double_t dX = (xax->GetXmax()-xax->GetXmin())/nbins;
  	
	while (i*step <= xax->GetXmax()-1) {
    	bin_index = xax->FindBin(i*step);
    	Xi = xax->GetXmin()+(bin_index+0.5)*dX;
    	sprintf(binlabel,"%.0f",Xi*ScaleFactor);
    	xax->SetBinLabel(bin_index,binlabel);
    	i+=1;
    }
}

//----------------------------------------------------------------------------------------------

//main
void ICSrootViewer()
{	
	gROOT->Reset();

	//----------------------------------input--------------------------------
  	//open an input file 
  	TFile* file = TFile::Open("output.root");
  	
  	Bool_t IsReducedfile = true;
  	
  	//define theta cuts
  	Double_t nbinstheta = 100;
  	Double_t thetaMax = 0.005;
  	Char_t cuttheta[128];
  	sprintf(cuttheta,"acos(momz/(momx^2+momy^2+momz^2)^0.5)<%.6f",thetaMax);
  	
  	//set Histogram parameters
  	Int_t nbinsE = 200; 		
  	Double_t Emin = 0;              //keV
  	Double_t Emax = 60; 			//keV
  	
  	Double_t binXsize = 1.; 		//um
  	Double_t binXStart = -50.;		//um
  	Double_t binXEnd = -binXStart; 
  	Int_t nbinsX = binXEnd*2./binXsize; 
  	
  	Double_t binYsize = 1.; 		//um   
  	Double_t binYStart = -50.;		//um
  	Double_t binYEnd = binXEnd; 
  	Int_t nbinsY = binYEnd*2./binYsize;  
  	
  	//options
  	Bool_t IWantSpectrum = false;
  	Bool_t IWantPhotonsVsTheta = false;
  	Bool_t IWantEVsTheta = false;
  	Bool_t IWantPlot2D = false;
  	 	 	  	
  	//disable statistics
  	gStyle->SetOptStat(0);
  	
  	//export  	
  	Bool_t IwantExport = true;
  	Char_t ExportFileName[128];
  	sprintf(ExportFileName,"output.dat");	
  	//----------------------------------------------------------------------- 	

  	
  	//original tree extraction
  	Float_t x,y,z,E,u,v,w,sx,sy,sz;
  	TTree* t1 = (TTree*)file->Get("ntuple");
  	t1->SetBranchAddress("posx",&x);
   	t1->SetBranchAddress("posy",&y);
    t1->SetBranchAddress("posz",&z);
  	t1->SetBranchAddress("e",&E);
   	t1->SetBranchAddress("momx",&u);
   	t1->SetBranchAddress("momy",&v);
    t1->SetBranchAddress("momz",&w);
    if (!IsReducedfile) {
		t1->SetBranchAddress("Sx",&sx);
	   	t1->SetBranchAddress("Sy",&sy);
		t1->SetBranchAddress("Sz",&sz);
    }
    Int_t N = t1->GetEntries();
    cout << endl << "Nph: " << N << endl << endl;
    const Int_t Nval = N;
       	
       	
    //photons VS theta
  	TH1D* Gammas_theta = new TH1D("Gammas_theta", "photons VS #theta", nbinstheta, 0, thetaMax);
  	t1->Project("Gammas_theta", "acos(momz/(momx^2+momy^2+momz^2)^0.5)", cuttheta); 
  	if (IWantPhotonsVsTheta) {			 
	  	Gammas_theta->SetLineColor(kRed);
	  	Gammas_theta->SetLineWidth(2);
	  	Gammas_theta->SetXTitle("#theta (rad)");
	  	Gammas_theta->GetXaxis()->CenterTitle();
	  	Gammas_theta->GetXaxis()->SetTitleOffset(1.2); 
	  	Gammas_theta->SetYTitle("Counts");
	  	Gammas_theta->GetYaxis()->CenterTitle();
	  	Gammas_theta->GetYaxis()->SetTitleOffset(1.5); 
	  	TCanvas* c1 = new TCanvas("c1","",1200,800); 
	  	c1->cd(); 									
	  	Gammas_theta->Draw("h");
  	}

     
    //Energy vs theta
    if (IWantEVsTheta) {	
		Double_t energy[Nval];
		Double_t theta[Nval];   
		for (Int_t i=0; i<Nval; i++){
			t1->GetEntry(i);  
			energy[i] = E*0.001;
			theta[i] = acos(w/pow(pow(u,2)+pow(v,2)+pow(w,2),0.5));
		}         
		TGraph* gr1 = new TGraph(Nval, theta, energy); 	
	   	gr1->SetTitle("Energy vs #theta");
	   	gr1->GetXaxis()->SetTitle("#theta (rad)");
	   	gr1->GetXaxis()->CenterTitle();
	   	gr1->GetXaxis()->SetTitleOffset(1.1); 
		gr1->GetYaxis()->SetTitle("E (keV)");
	   	gr1->GetYaxis()->CenterTitle();
	  	gr1->GetYaxis()->SetTitleOffset(1.2); 
	  	gr1->GetXaxis()->SetRangeUser(0.0,0.1);
	  	TCanvas* c2 = new TCanvas("c2","",1200,800); 
		c2->cd();
	   	gr1->Draw("ap");     
   	}
   	
  	
   	//photon spectrum
   	if (IWantSpectrum) {
	  	TH1D* Spectrum = new TH1D("Spectrum", "", nbinsE, Emin, Emax);
	  	t1->Project("Spectrum", "e*0.001");  	 
	  	Spectrum->SetLineColor(kBlack);
	  	Spectrum->SetLineWidth(2); 
	  	Spectrum->SetXTitle("Energy (keV)");
	  	Spectrum->GetXaxis()->CenterTitle();
	  	Spectrum->GetXaxis()->SetTitleOffset(1.2); 
	  	Spectrum->SetYTitle("Counts");
	  	Spectrum->GetYaxis()->CenterTitle();
	  	Spectrum->GetYaxis()->SetTitleOffset(1.3); 	 
	  	TCanvas* c3 = new TCanvas("c3","",1200,800); 	 	
	  	c3->cd(); 
	  	Spectrum->Draw("h"); 
	  	TH1D* collSpectrum = new TH1D("collSpectrum", "", nbinsE, Emin, Emax);
	  	t1->Project("collSpectrum", "e*0.001", cuttheta); 		
	  	collSpectrum->SetLineColor(kRed);
	  	collSpectrum->SetLineWidth(2); 
	  	collSpectrum->Draw("same"); 	
	  	
	  	Int_t binpeak = collSpectrum->GetMaximumBin(); 
	  	Double_t Epeak = collSpectrum->GetXaxis()->GetBinCenter(binpeak); 	
	  	cout << "peak energy: " << Epeak << " keV" << endl;
	  	cout << "mean energy: " << collSpectrum->GetMean() << " keV" << endl;
	  	cout << "BW of collimated (" << thetaMax*1000 << " mrad) spectrum: "
		     << collSpectrum->GetRMS()*100/collSpectrum->GetMean() << "%" 
		     << endl << endl;
    }
  	
  	
  	//spatial distribution
	TH2D* SpatialDistribution = new TH2D("SpatialDistribution", "", nbinsX,
																    binXStart, 
																    binXEnd, 
																    nbinsY, 
																    binYStart, 
																    binYEnd); 
	t1->Project("SpatialDistribution","posy*1000000:posx*1000000"); 
	
	TProfile2D* Eprofile = new TProfile2D("Eprofile", "", nbinsX, 
	                                                      binXStart, 
	                                                      binXEnd, 
	                                                      nbinsY, 
	                                                      binYStart, 
	                                                      binYEnd);	
	
	if (IWantPlot2D) {
		SpatialDistribution->SetXTitle("x (um)");
		SpatialDistribution->GetXaxis()->CenterTitle();
		SpatialDistribution->GetXaxis()->SetTitleOffset(1.2); 
		SpatialDistribution->SetYTitle("y (um)");
		SpatialDistribution->GetYaxis()->CenterTitle();
		SpatialDistribution->GetYaxis()->SetTitleOffset(1.3); 
		TCanvas* c4 = new TCanvas("c4","",1600,1100); 
		c4->Divide(2,1); 
		c4->cd(1);
		SpatialDistribution->SetContour(50);
		SpatialDistribution->Draw("colz");
		
		for (Int_t i=0; i<N; i++) {
   			t1->GetEntry(i);
      		Eprofile->Fill(x*1.e6,y*1.e6,E*0.001,1);
   		} 
   		c4->cd(2);
	   	Eprofile->Draw("cont4z");
		Eprofile->SetContour(50);
		Eprofile->SetXTitle("x (um)");
	  	Eprofile->GetXaxis()->CenterTitle();
	  	Eprofile->GetXaxis()->SetTitleOffset(1.1); 
	  	Eprofile->SetYTitle("y (um)");
	  	Eprofile->GetYaxis()->CenterTitle();
	  	Eprofile->GetYaxis()->SetTitleOffset(1.3); 
	  	Eprofile->GetZaxis()->SetRangeUser(Emin, Emax);
	}
  	
  	
  	//export
  	if (IwantExport) {
  		//open a text file
  		ofstream f(ExportFileName); 
    	if (!f) {
        	cout << "Error opening the file!";
        	return;
    	}
    	for (Int_t i=0; i<Nval; i++) {
    		t1->GetEntry(i); 
			if (IsReducedfile) {
				f << x << " " 
				  << y << " " 
				  << z << " " 
				  << E << " " 
				  << u << " " 
				  << v << " " 
				  << w << endl;	
		    } else {
				f << x << " " 
				  << y << " " 
				  << z << " " 
				  << E << " " 
				  << u << " " 
				  << v << " " 
				  << w << " " 
				  << sy << " " 
				  << sz << " " 
				  << sx << endl;	
				  //NOTE: for Geant4 simulations, si are Stokes parameters defined 
				  //differently from CAIN (check required - sept 2021)
			}
    	} 
  		//close the text file
    	f.close(); 
   		cout << "Export of the root file to a text file successful!" << endl << endl;
   	}

}