#include <TROOT.h>
#include <TH1.h>
#include <TFile.h>
#include <TTree.h>
#include <TCanvas.h>
#include <TStyle.h>
#include <TColor.h>
#include <TGraph.h>
#include <TGraphErrors.h>
#include <TF1.h>
#include <TLatex.h>
#include <TProfile.h>

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


// The file tavg_res.cc is created by the tavg_check_tgraph.C macro
#include "tavg_res.cc"
#include "tavg.cc"
#include "E_calibrate.cc"

#include "Angle.h"
#include "Geom.h"

#include "StandardLabels.C"

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

//------------------------------------
// Copied from JEventProcessor_bcal_timing.cc
typedef struct{
	int event;
	int layer;
	int sector;
	int fADC_up;
	int fADC_dn;
	float Etot;
	float geometric_mean;
	float tup;
	float tdn;
	float tup_corrected;
	float tdn_corrected;
	float theta_thrown;
	float E_thrown;
}HIT_t;
//------------------------------------

//-----------------
// tavg_res_vs_E
//-----------------
void tavg_res_vs_E(void)
{
	gStyle->SetStatY(0.90);
	gStyle->SetOptStat(2200);
	TColor::CreateColorWheel();

	TFile *f = new TFile("hd_root.root");
	f->cd();
	
	TTree *tree = (TTree*)gROOT->FindObject("tree");
	HIT_t hit;
	HIT_t *hitptr = &hit;
	TBranch *branch = tree->GetBranch("T");
	branch->SetAddress(hitptr);
	
	TH2D *h = new TH2D("h", "", 200, 0.03, 2.5, 60, -600.0, 600.0);
	
	// Loop over all events and fill histogram. We do this
	// rather than TTree::Project since we need to call the 
	// BCAL_tavg_res for each entry
	int Nentries = branch->GetEntries();
	double t_tot = 0.0;
	double A_tot = 0.0;
	double E_tot = 0.0;
	int last_event=1;
	for(int i=1; i<=Nentries; i++){
		tree->GetEntry(i);
		
		// Fill histogram on event boundaries and reset for next event
		if(hit.event!=last_event){

			t_tot /= A_tot;
			
			E_tot = BCAL_fADC2E(E_tot); // convert to GeV

			h->Fill(E_tot, t_tot);
			
			t_tot = 0.0;
			A_tot = 0.0;
			E_tot = 0.0;
			last_event = hit.event;
		}
		
		if(hit.layer<1 || hit.layer>MaxSummedLayers())continue;
		if(hit.tup_corrected<-100)continue;
		if(hit.tdn_corrected<-100)continue;
		
		// Accumulate info for this event
		double geomn = hit.geometric_mean;
		double E = geomn; // convert to GeV just before filling histo above
		double weight = E;
		double tavg_corrected = (hit.tup_corrected + hit.tdn_corrected)/2.0 - BCAL_tavg(geomn, hit.layer);
		if(!IsOuter(hit.layer)){
			t_tot += weight*1000.0*tavg_corrected;
			A_tot += weight;
		}
		E_tot += E;
	}
	
	//h->FitSlicesY();
	//TH1D *h_2 = (TH1D*)gROOT->FindObject("h_2");
	
	TGraphErrors *gres;
	TGraph *gmean;
	GetPoints(h, gres, gmean);
	
	// Create a fit function
	TF1 *fun = new TF1("fun", "sqrt([0]*[0] + pow([1]/sqrt(x) ,2.0))");
	fun->SetParameter(0, 10.0);
	fun->SetParameter(1, 50.0);
	fun->SetLineColor(kRed);
	fun->SetLineWidth(2);
	
	//g->Fit(fun, "", "", 0.01, 2.4);
	
	fun->SetLineColor(kBlue);
	gres->Fit(fun,"","", 0.0,2.0);

	TCanvas *c1 = new TCanvas("c1");
	c1->SetTicks();
	c1->SetGrid();
	
	TH2D *axes = new TH2D("axes","", 100, 0.0, 2.5, 100, 0.0, 400.0);
	axes->SetStats(0);
	axes->SetXTitle("Geometric mean (GeV)");
	axes->SetYTitle("Energy weighted t_{avg} #sigma (ps)");
	axes->Draw();
	
	gres->Draw("Psame");
	
	char str[256];
	sprintf(str, "#sigma_{tavg} = #frac{%4.1f ps}{#sqrt{E}} #oplus %4.1f ps", fun->GetParameter(1), fun->GetParameter(0));
	TLatex *lab = new TLatex(0.4, 250, str);
	lab->SetTextSize(0.05);
	lab->Draw();

	StandardLabels(axes, Scheme(), "", AngleStr("#theta_{#gamma}="));
	
	c1->SaveAs("tavg_res_vs_E.png");
	c1->SaveAs("tavg_res_vs_E.pdf");
}

//------------------
// GetPoints
//------------------
void GetPoints(TH2D *h, TGraphErrors* &gres, TGraph* &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;
	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()>=20){
		
			// Fit to Gaussian
			h_py->Fit("gaus","0Q");
			//c1->Update();

			// Get fit results
			double mean = h_py->GetFunction("gaus")->GetParameter(1);
			double sigma = h_py->GetFunction("gaus")->GetParameter(2);
			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 = 5.0; // ps
			sigma_err = sqrt(sigma_err*sigma_err + epsilon*epsilon);
			
			// Error on reconstructed energy is difficult to estimate here.
			// It is definitely worse than the final algorithm will yield.
			// The fact that the "reconstructed" energy extends out to 2.5GeV
			// when the generated data set only went to 2GeV says it is 
			// considerably worse than 5.5%/sqrt(E). For lack of a better
			// idea, I use the RMS of the fADC values about the mean
			// in the current range.
			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);
			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->SetMarkerStyle(20);
	gres->SetMarkerSize(1.0);
	gres->SetLineColor(kMagenta);
	gres->SetLineWidth(2.0);

	gmean = new TGraph(Npoints, &x[0], &y_mean[0]);
	gmean->SetMarkerColor(kRed);
	gmean->SetMarkerStyle(22);
	gmean->SetLineColor(kRed);
}