// This macro will make a ROOT file containing histograms
// of the charged tracking resolutions suitable for use
// with the src/programs/Simulation/plugins/hdparsim
// plugin for doing parameteric simulations

#include "filelist.h"

//--------------------
// mk_hd_photon_res
//--------------------
void mk_hd_photon_res(void)
{
	gROOT->Reset();
	gStyle->SetPalette(1);

	// Create chain of all the input trees
	TChain *phtneff = new TChain("CALORIMETRY/phtneff","phtneff");
	AddFilesToChain(phtneff);
	
	// Create a 2D histo with the binning and limits that will be
	// used for all histograms. We want all resolutions and efficiencies
	// to be a function of total momentum and theta angle
	TH2D *axes = new TH2D("axes","nada", 200, 0.0, 180.0, 100, 0.0, 8.0);
	axes->SetXTitle("Polar angle #theta (degrees)");
	axes->SetYTitle("Total momentum (GeV/c)");
	axes->SetStats(0);
	
	// Standard cut
	TCut cut1("abs((E.pfit.Mag()-E.pthrown.Mag())/E.pthrown.Mag())<0.20");
	
	// Open output file to hold histograms
	TFile *f = new TFile("hd_res_photon.root", "RECREATE");
	
	// Create resolution and efficiency histograms
	MakeEfficiency(phtneff, axes, cut1);
	MakeERes(phtneff, axes, cut1);
	MakeThetaRes(phtneff, axes, cut1);
	MakePhiRes(phtneff, axes, cut1);

	// Close resolutions file
	f->Write();
	//delete f;
}

//--------------------
// MakeEfficiency
//--------------------
void MakeEfficiency(TChain *phtneff, TH2D *axes, TCut &cut1)
{
	TH2D *h = axes->Clone("eff_vs_p_vs_theta");
	h->SetTitle("Photon Reconstruction Efficiency vs. total momentum and #theta angle");
	cout<<"Making "<<h->GetName()<<endl;

	TH2D *numerator = axes->Clone("numerator");
	TH2D *denominator = axes->Clone("denominator");
	
	cout<<"   Projecting numerator ..."<<endl;
	phtneff->Project("numerator","E.pthrown.Mag():E.pthrown.Theta()*TMath::RadToDeg()", cut1);
	cout<<"   Projecting denominator ..."<<endl;
	phtneff->Project("denominator","E.pthrown.Mag():E.pthrown.Theta()*TMath::RadToDeg()");

	cout<<"   Finalizing ..."<<endl;
	h->Divide(numerator, denominator);
	h->GetZaxis()->SetRangeUser(0.2, 1.00);
	h->Draw("colz");
	
	//delete numerator;
	//delete denominator;
}

//--------------------
// MakeERes
//--------------------
void MakeERes(TChain *phtneff, TH2D *axes, TCut &cut1)
{
	TH2D *h = axes->Clone("dE_over_E_vs_p_vs_theta");
	h->SetTitle("Relative Energy resolution vs. total momentum and #theta angle");
	cout<<"Making "<<h->GetName()<<endl;
	
	// Make 3D histo to project onto
	TAxis *xaxis = axes->GetXaxis();
	TAxis *yaxis = axes->GetYaxis();
	TH3D *tmp3d = new TH3D("tmp3d", ""
		, xaxis->GetNbins(), xaxis->GetXmin(), xaxis->GetXmax()
		, yaxis->GetNbins(), yaxis->GetXmin(), yaxis->GetXmax()
		, 50, -0.5, 0.5);
	
	cout<<"   Projecting  ..."<<endl;
	phtneff->Project("tmp3d","(E.pfit.Mag()-E.pthrown.Mag())/E.pthrown.Mag():E.pthrown.Mag():E.pthrown.Theta()*TMath::RadToDeg()", cut1);
	cout<<"   Fitting ..."<<endl;
	tmp3d->FitSlicesZ();

	cout<<"   Finalizing ..."<<endl;

	// Find histos resulting from fits
	TH2D *tmp3d_0 = (TH2D*)gROOT->FindObject("tmp3d_0");
	TH2D *tmp3d_1 = (TH2D*)gROOT->FindObject("tmp3d_1");
	TH2D *tmp3d_2 = (TH2D*)gROOT->FindObject("tmp3d_2");
	TH2D *tmp3d_chi2 = (TH2D*)gROOT->FindObject("tmp3d_chi2");
	
	// Copy the contents of sigma histo into h
	h->Add(tmp3d_2);
	h->GetZaxis()->SetRangeUser(0.0, 0.2);
	h->Draw("colz");

	h = (TH2D*)axes->Clone("Nevents");
	phtneff->Project("Nevents","E.pthrown.Mag():E.pthrown.Theta()*TMath::RadToDeg()", cut1);
	//return;	

	// Clean up
	delete tmp3d_0;
	delete tmp3d_1;
	delete tmp3d_2;
	delete tmp3d_chi2;
	delete tmp3d;
}

//--------------------
// MakeThetaRes
//--------------------
void MakeThetaRes(TChain *phtneff, TH2D *axes, TCut &cut1)
{
	TH2D *h = axes->Clone("dtheta_vs_p_vs_theta");
	h->SetTitle("Polar angular resolution vs. total momentum and #theta angle");
	TAxis *xaxis = axes->GetXaxis();
	TAxis *yaxis = axes->GetYaxis();
	cout<<"Making "<<h->GetName()<<endl;
	
	// It turns out there is enough variation in the width of the spectra between the
	// "good" resolutions parts and the "bad" resolution parts that we can't have a single
	// z-axis definition that allows good fits everywhere. Hence, we have to do this in
	// sections which takes considerably longer.
	int Nsections=5;
	double xmin[]={ 0.0, 20.0,  20.0,	5.0,  20.0};
	double xmax[]={20.0,  0.0,   0.0,  25.0,  40.0};
	double ymin[]={ 0.0,  2.0,   0.0,   0.9,   3.0};
	double ymax[]={ 0.0,  0.0,   2.0,   2.2,   0.0};
	double  lim[]={15.0, 50.0, 100.0,  50.0,  35.0};

	for(int i=0; i<Nsections; i++){

		// Make 3D histo to project onto
		TH3D *tmp3d = new TH3D("tmp3d", ""
			, xaxis->GetNbins(), xaxis->GetXmin(), xaxis->GetXmax()
			, yaxis->GetNbins(), yaxis->GetXmin(), yaxis->GetXmax()
			, 30, -lim[i], lim[i]);
		
		cout<<"   Projecting  "<<tmp3d->GetName()<<" for "<<i<<"..."<<endl;
		phtneff->Project(tmp3d->GetName(),"(E.pfit.Theta()-E.pthrown.Theta())*1000.0:E.pthrown.Mag():E.pthrown.Theta()*TMath::RadToDeg()", cut1);

		int xbinmin = xaxis->FindBin(xmin[i]);
		int xbinmax = xaxis->FindBin(xmax[i]!=0.0 ? xmax[i]:xaxis->GetNbins());
		int ybinmin = yaxis->FindBin(ymin[i]);
		int ybinmax = yaxis->FindBin(ymax[i]!=0.0 ? ymax[i]:yaxis->GetNbins());
		cout<<"   Fitting "<<tmp3d->GetName()<<"..."<<endl;
		tmp3d->FitSlicesZ(0, xbinmin, xbinmax, ybinmin, ybinmax);
		
		// Find pointers to resulting histos
		TH2D *tmp3d_0 = (TH2D*)gROOT->FindObject("tmp3d_0");
		TH2D *tmp3d_1 = (TH2D*)gROOT->FindObject("tmp3d_1");
		TH2D *tmp3d_2 = (TH2D*)gROOT->FindObject("tmp3d_2");
		TH2D *tmp3d_chi2 = (TH2D*)gROOT->FindObject("tmp3d_chi2");

		// Copy results into "h"
		cout<<"   Copying fit results "<<tmp3d->GetName()<<"..."<<endl;
		for(int ix=xbinmin; ix<=xbinmax; ix++){
			for(int iy=ybinmin; iy<=ybinmax; iy++)_{
				h->SetBinContent(ix, iy, tmp3d_2->GetBinContent(ix, iy));
			}
		}

		// clean up
		delete tmp3d_0;
		delete tmp3d_1;
		delete tmp3d_2;
		delete tmp3d_chi2;
		delete tmp3d;
	}

	// Some fits just fail and leave a bin content of 0. Check every bin and
	// if it is zero, use the value from the bin to its left
	cout<<"   Finalizing ..."<<endl;
	for(int ix=2; ix<=xaxis->GetNbins(); ix++){
		if(xaxis->GetBinCenter(ix)>125.0)continue;
		for(int iy=1; iy<=yaxis->GetNbins(); iy++){
			if(h->GetBinContent(ix,iy)==0.0){
				h->SetBinContent(ix,iy, h->GetBinContent(ix-1,iy));
			}
		}
	}

	// Copy the contents of sigma histo into h
	h->GetZaxis()->SetRangeUser(0.0, 200.0);
	h->Draw("colz");
}

//--------------------
// MakePhiRes
//--------------------
void MakePhiRes(TChain *phtneff, TH2D *axes, TCut &cut1)
{
	TH2D *h = axes->Clone("dphi_vs_p_vs_theta");
	h->SetTitle("Azimuthal angular resolution vs. total momentum and #theta angle");
	cout<<"Making "<<h->GetName()<<endl;
	
	// Make 3D histo to project onto
	TAxis *xaxis = axes->GetXaxis();
	TAxis *yaxis = axes->GetYaxis();
	TH3D *tmp3d = new TH3D("tmp3d", ""
		, xaxis->GetNbins(), xaxis->GetXmin(), xaxis->GetXmax()
		, yaxis->GetNbins(), yaxis->GetXmin(), yaxis->GetXmax()
		, 30, -30.0, 30.0);
	
	cout<<"   Projecting  ..."<<endl;
	phtneff->Project("tmp3d","(E.pfit.Phi()-E.pthrown.Phi())*1000.0:E.pthrown.Mag():E.pthrown.Theta()*TMath::RadToDeg()", cut1);
	cout<<"   Fitting ..."<<endl;
	tmp3d->FitSlicesZ();

	cout<<"   Finalizing ..."<<endl;

	// Find histos resulting from fits
	TH2D *tmp3d_0 = (TH2D*)gROOT->FindObject("tmp3d_0");
	TH2D *tmp3d_1 = (TH2D*)gROOT->FindObject("tmp3d_1");
	TH2D *tmp3d_2 = (TH2D*)gROOT->FindObject("tmp3d_2");
	TH2D *tmp3d_chi2 = (TH2D*)gROOT->FindObject("tmp3d_chi2");
	
	// Copy the contents of sigma histo into h
	h->Add(tmp3d_2);
	h->GetZaxis()->SetRangeUser(0.0, 200.0);
	h->Draw("colz");
	//return;	

	// Clean up
	delete tmp3d_0;
	delete tmp3d_1;
	delete tmp3d_2;
	delete tmp3d_chi2;
	delete tmp3d;
}