// This will fit histograms from the tree in hd_root.root to
// obtain a set of timewalk corrections functions.
//
// This is done by projecting the uncorrected times onto a
// histogram in bins of fADC value. The maximum is then used to
// fill a histo and fit it to a specific functional form and code is
// generated to capture those fit results into a C++ accessible
// format: timewalk_max.cc.

#include <iostream>
#include <fstream>
using namespace std;

#include <TROOT.h>
#include <TFile.h>
#include <TH2.h>
#include <TF1.h>
#include <TCanvas.h>
#include <TTree.h>
#include <TLegend.h>
#include <TGraphErrors.h>

#include "Geom.h"
#include "Angle.h"
#include "StandardLabels.C"

void GetPoints(TH2D *h, TGraphErrors* &gres, TGraphErrors* &gmean);

void timewalk_calibrate(void)
{

	TFile *f = new TFile("hd_root.root");
	f->cd(); // avoid compiler warnings when running with "+"
	
	TTree *tree = (TTree*)gROOT->FindObject("tree");

	TCanvas *c1 = new TCanvas("c1");
	c1->SetTicks();
	c1->SetGrid();
	
	TH2D *axes = new TH2D("axes","", 100, 0.0, 800.0, 100, Tmin()-5.0, Tmax()+15.0);
	axes->SetStats(0);
	axes->SetXTitle("fADC");
	axes->SetYTitle("t (ns)");

	TF1 *fun_up = new TF1("fun_up", "[0] + [1]/(pow(x,[2])+[3])", 0.0, 800.0);
	TF1 *fun_dn = new TF1("fun_dn", "[0] + [1]/(pow(x,[2])+[3])", 0.0, 800.0);
	fun_up->SetLineWidth(2);
	fun_dn->SetLineWidth(2);
	
	// Estimated values for tup and tdn
	double tup_est, tdn_est;
	GetTupTdn(tup_est, tdn_est, (17.0+407.0)/2.0);
	
	fun_up->SetLineColor(kBlue);
	fun_dn->SetLineColor(kRed);
	
	double parms_up[10+1][4];
	double parms_dn[10+1][4];

	// Open PDF file for plots
	c1->Print("timewalk.ps[");

	// Loop over layers
	int Nlayers=0;
	for(int ilayer=1; ilayer<=MaxSummedLayers() ; ilayer++, Nlayers++){
	//for(int ilayer=2; ilayer<=2 ; ilayer++, Nlayers++){

		// First, project onto a 2D histo
		TH2D *hup_2D = new TH2D("hup_2D","", 800, 40.0, 800.0, 600, 0.0, 60.0);
		TH2D *hdn_2D = new TH2D("hdn_2D","", 800, 40.0, 800.0, 600, 0.0, 60.0);
		
		char cut[256];
		sprintf(cut, "layer==%d", ilayer);
		tree->Project("hup_2D", "tup:fADC_up", cut);
		tree->Project("hdn_2D", "tdn:fADC_dn", cut);
		
		// Create TGraphs from histos
		TGraphErrors *gres_up, *gmean_up;
		TGraphErrors *gres_dn, *gmean_dn;
		GetPoints(hup_2D, gres_up, gmean_up);
		GetPoints(hdn_2D, gres_dn, gmean_dn);
	
		gmean_up->SetLineColor(kCyan);
		gmean_dn->SetLineColor(kMagenta);
		gmean_up->SetMarkerColor(kCyan);
		gmean_dn->SetMarkerColor(kMagenta);
		gmean_up->SetLineWidth(2);
		gmean_dn->SetLineWidth(2);

		char title[256];
		sprintf(title, "Timewalk for layer %d", ilayer);
		axes->SetTitle(title); 

		cout<<"tup_est="<<tup_est<<" tdn_est="<<tdn_est<<endl;
		
		fun_up->SetParameter(0, tup_est);
		fun_up->SetParameter(1, 0.04);
		fun_up->SetParameter(2, 0.006);
		fun_up->SetParameter(3, -1.0);

		fun_dn->SetParameter(0, tdn_est);
		fun_dn->SetParameter(1, 0.04);
		fun_dn->SetParameter(2, 0.006);
		fun_dn->SetParameter(3, -1.0);
		
		double xmin_up, xmin_dn, tmp;
		gmean_up->GetPoint(0, xmin_up, tmp);
		gmean_dn->GetPoint(0, xmin_dn, tmp);

		gmean_up->Fit(fun_up,"","",xmin_up, 800.0);
		for(int ipar=0; ipar<=3; ipar++)parms_up[ilayer][ipar] = fun_up->GetParameter(ipar);

		gmean_dn->Fit(fun_dn,"","",xmin_dn, 800.0);
		for(int ipar=0; ipar<=3; ipar++)parms_dn[ilayer][ipar] = fun_dn->GetParameter(ipar);

		axes->Draw();
		tree->Draw("tup:fADC_up", cut, "same");
		tree->Draw("tdn:fADC_dn", cut, "same");
		gmean_up->Draw("Psame");
		gmean_dn->Draw("Psame");

		TLegend *leg = new TLegend(0.6, 0.70, 0.8, 0.85);
		leg->SetFillColor(kWhite);
		leg->AddEntry(fun_up, "upstream");
		leg->AddEntry(fun_dn, "downstream");
		leg->Draw();

		StandardLabels(axes, Scheme(), AngleStr("#theta_{#gamma}="));

		c1->Update();
		c1->Print("timewalk.ps");
		
		char fname[256];
		sprintf(fname, "timewalk_calibrate_layer%0d.png", ilayer);
		c1->SaveAs(fname);
	}
	
	// Close PDF file 
	c1->Print("timewalk.ps]");

	// Write parameters out in form of C++ file
	char cppfname[256];
	sprintf(cppfname, "timewalk_%s.h", Scheme());
	ofstream pout(cppfname);
	pout<<endl;
	pout<<"#include<cmath>"<<endl;
	pout<<"using namespace std;"<<endl;
	pout<<endl;
	pout<<"#define NLAYERS_BCAL "<<Nlayers<<endl;
	pout<<"double BCAL_TimeWalkUp(double fADC, int layer)"<<endl;
	pout<<"{"<<endl;
	pout<<"	double tw_par[NLAYERS_BCAL][4] ={"<<endl;
	for(int i=1; i<=Nlayers; i++){
		pout<<"		{";
		for(int j=0; j<=3; j++){
			if(j!=0)pout<<", ";
			pout<<parms_up[i][j];
		}
		pout<<"}";
		if(i<Nlayers)pout<<",";
		pout<<"\t// layer "<<i<<endl;
	}
	pout<<"		};"<<endl;
	pout<<endl;
	pout<<"	layer = layer<1 ? 1:layer>NLAYERS_BCAL ? NLAYERS_BCAL:layer;"<<endl;
	pout<<"	double *p = tw_par[layer-1];"<<endl;
	pout<<"	return p[0]+p[1]/(pow(fADC,p[2])+p[3]);"<<endl;
	pout<<"}"<<endl;
	pout<<endl;
	pout<<endl;
	pout<<"double BCAL_TimeWalkDn(double fADC, int layer)"<<endl;
	pout<<"{"<<endl;
	pout<<"	double tw_par[NLAYERS_BCAL][4] ={"<<endl;
	for(int i=1; i<=Nlayers; i++){
		pout<<"		{";
		for(int j=0; j<=3; j++){
			if(j!=0)pout<<", ";
			pout<<parms_dn[i][j];
		}
		pout<<"}";
		if(i<Nlayers)pout<<",";
		pout<<"\t// layer "<<i<<endl;
	}
	pout<<"		};"<<endl;
	pout<<endl;
	pout<<"	layer = layer<1 ? 1:layer>NLAYERS_BCAL ? NLAYERS_BCAL:layer;"<<endl;
	pout<<"	double *p = tw_par[layer-1];"<<endl;
	pout<<"	return p[0]+p[1]/(pow(fADC,p[2])+p[3]);"<<endl;
	pout<<"}"<<endl;
	pout.close();
	cout<<endl;
	cout<<"Wrote timewalk parameters to \""<<cppfname<<"\""<<endl;
	cout<<endl;
	
	cout<<"Plots written to \"timewalk.ps\". Convert to PDF with:"<<endl;
	cout<<endl;
	cout<<"  ps2pdf timewalk.ps"<<endl;
	cout<<endl;
	cout<<"attempting to run for you ...."<<endl;
	system("ps2pdf timewalk.ps");
	cout<<"Done."<<endl;
}


//------------------
// GetPoints
//------------------
void GetPoints(TH2D *h, TGraphErrors* &gres, TGraphErrors* &gmean)
{
	int min_entries = 300;

	// Project onto the x-axis so we can find the limits for
	// for the projection on the y-axis based on number of events
	TH1D *h_px = h->ProjectionX("_px", 1);

	// Loop over (uneven) bins
	int Npoints = 0;
	vector<double> x;
	vector<double> y;
	vector<double> y_mean;
	vector<double> xerr;
	vector<double> yerr;
	vector<double> y_mean_err;
	int Nbins = h->GetNbinsX();
	for(int start_bin=1; start_bin<=Nbins; ){

		// Find limits to give us min_entries entries
		int end_bin = start_bin;
		while(h_px->Integral(start_bin, end_bin)<min_entries && end_bin<Nbins){
			++end_bin;
		}
		
		// If only a few events are left, go ahead and include them here.
		if(end_bin<Nbins && h_px->Integral(end_bin, Nbins)<min_entries)end_bin=Nbins;

		// Project onto y-axis
		cout<<"projecting bins "<<start_bin<<" to "<<end_bin<<"  Nentries="<<h_px->Integral(start_bin, end_bin)<<endl;
		TH1D *h_py = h->ProjectionY("_py",start_bin, end_bin);

		if(h_py->Integral()>=5){
		
			// Fit to Gaussian
			h_py->Fit("gaus","0Q");
			//c1->Update();

			// Get fit results
			//double mean = h_py->GetFunction("gaus")->GetParameter(1);
double mean = h_py->GetMean();
			double sigma = h_py->GetFunction("gaus")->GetParameter(2);
			//double mean_err = h_py->GetFunction("gaus")->GetParError(1);
double mean_err = h_py->GetRMS();
			double sigma_err = h_py->GetFunction("gaus")->GetParError(2);

			// Find average fADC value for hits in this bin range
			double fADC = 0.0;
			double norm = 0.0;
			for(int bin=start_bin; bin<=end_bin; bin++){
				double weight = h_px->GetBinContent(bin);
				fADC += h_px->GetBinCenter(bin)*weight;
				norm += weight;
			}
			fADC /= norm;
		
			// Add floor term to error obtained from Gaussian fit
			double epsilon = 0.200; // ns
			mean_err = sqrt(mean_err*mean_err + epsilon*epsilon);
			sigma_err = sqrt(sigma_err*sigma_err + epsilon*epsilon);

			// Estimate error in x
			double fADC_rms = 0.0;
			double norm_rms = 0.0;
			for(int bin=start_bin; bin<=end_bin; bin++){
				double weight = h_px->GetBinContent(bin);
				fADC_rms += pow(h_px->GetBinCenter(bin)-fADC, 2.0)*weight;
				norm_rms += weight;
			}
			fADC_rms /= norm_rms;
			fADC_rms = sqrt(fADC_rms);

			x.push_back(fADC);
			y.push_back(sigma);
			y_mean.push_back(mean);
			xerr.push_back(fADC_rms);
			yerr.push_back(sigma_err);
			y_mean_err.push_back(mean_err);
			Npoints++;
		}

		// Increment for next iteration
		start_bin=end_bin+1;
	}
	
	gres = new TGraphErrors(Npoints, &x[0], &y[0], &xerr[0], &yerr[0]);
	gres->SetMarkerColor(kMagenta);
	gres->SetLineColor(kMagenta);
	gres->SetLineWidth(2);

	gmean = new TGraphErrors(Npoints, &x[0], &y_mean[0], &xerr[0], &y_mean_err[0]);
	gmean->SetMarkerColor(kRed);
	gmean->SetMarkerStyle(23);
	gmean->SetLineColor(kRed);
	
}