// When the hd_res_charged.root file is made, occasionally there is 
// a bin in a resolution histogram where the fit fails and the bin
// is left with a non-sensical value. This macro will look through 
// all of the bin contents and when it finds such a bin it will copy
// the contents of the bin immediately to it's left.
//
// The input file is specified as an argument to this macro
// (should be something like "hd_res_charged.root") and the
// output will be the same filename but with "touched_up_"
// prepended.

//----------------
// touch_up_hd_res_charged
//----------------
void touch_up_hd_res_charged(const char *fname)
{
	gROOT->Reset();

	// Open input file and get pointers to all histograms
	TFile *ifile = new TFile(fname);
	if(!ifile->IsOpen()){
		cout<<"Unable to open input file \""<<fname<<"\"!"<<endl;
		return;
	}

	TH2D *eff_vs_p_vs_theta						= (TH2D*)gROOT->FindObject("eff_vs_p_vs_theta");
	TH2D *eff_numerator							= (TH2D*)gROOT->FindObject("eff_numerator");
	TH2D *eff_denominator						= (TH2D*)gROOT->FindObject("eff_denominator");
	TH2D *dp_over_p_vs_p_vs_theta				= (TH2D*)gROOT->FindObject("dp_over_p_vs_p_vs_theta");
	TH2D *dpt_over_pt_vs_p_vs_theta			= (TH2D*)gROOT->FindObject("dpt_over_pt_vs_p_vs_theta");
	TH2D *dtheta_vs_p_vs_theta					= (TH2D*)gROOT->FindObject("dtheta_vs_p_vs_theta");
	TH2D *dphi_vs_p_vs_theta					=	(TH2D*)gROOT->FindObject("dphi_vs_p_vs_theta");
	TH3D *trk_chisq_per_Ndof_vs_p_vs_theta	= (TH3D*)gROOT->FindObject("trk_chisq_per_Ndof_vs_p_vs_theta");
	TH2D *Ncdchits_vs_p_vs_theta				= (TH2D*)gROOT->FindObject("Ncdchits_vs_p_vs_theta");
	TH2D *Nfdchits_vs_p_vs_theta				= (TH2D*)gROOT->FindObject("Nfdchits_vs_p_vs_theta");

	// Open output file
	char ofname[256];
	sprintf(ofname,"touched_up_%s",fname);
	TFile *ofile = new TFile(ofname, "RECREATE");
	
	// Touch up histos
	touch_up(eff_vs_p_vs_theta, 1.0E-6, 1.0);
	touch_up(dp_over_p_vs_p_vs_theta, 1.0E-4, 10.0);
	touch_up(dpt_over_pt_vs_p_vs_theta, 1.0E-4, 10.0);
	touch_up(dtheta_vs_p_vs_theta, 1.0E-2, 1000.0);
	touch_up(dphi_vs_p_vs_theta, 1.0E-2, 1000.0);

	// Copy histos into output file
	eff_vs_p_vs_theta->Clone();
	eff_numerator->Clone();
	eff_denominator->Clone();
	dp_over_p_vs_p_vs_theta->Clone();
	dpt_over_pt_vs_p_vs_theta->Clone();
	dtheta_vs_p_vs_theta->Clone();
	dphi_vs_p_vs_theta->Clone();
	trk_chisq_per_Ndof_vs_p_vs_theta->Clone();
	Ncdchits_vs_p_vs_theta->Clone();
	Nfdchits_vs_p_vs_theta->Clone();
	
	// Close output file
	ofile->Write();
	delete ofile;
	delete ifile;
}

//----------------
// touch_up
//----------------
void touch_up(TH2D *h, double zmin, double zmax)
{
	int Nbins_touched=0;

	TAxis *xaxis = h->GetXaxis();
	TAxis *yaxis = h->GetYaxis();

	// Bleed values to the right for out of range bins
	for(int iy=1; iy<=yaxis->GetNbins(); iy++){
	
		// The right-hand side of the plot will naturally have zeros due to 
		// geometrical acceptance. We don't want to bleed into that region
		// so we need to find the limit where to stop bleeding.
		int xbin_max;
		for(xbin_max=xaxis->GetNbins(); xbin_max>2; xbin_max--){
			double z = h->GetBinContent(xbin_max,iy);
			if(z>=zmin && z<=zmax)break;
		}

		for(int ix=2; ix<xbin_max; ix++){
			double z = h->GetBinContent(ix,iy);
			if(z<zmin || z>zmax){
				h->SetBinContent(ix,iy, h->GetBinContent(ix-1,iy));
				Nbins_touched++;
			}
		}
	}
	
	// Some bins aren't out of range, but are inconsistent with their neighbors indicating
	// something is wrong. Loop over each bin and compare it to the average of its neighbors
	// to check for consistency. We check the next to nearest neighbors since it seems these
	// problems ofter occur in neighboring bins.
	// Bleed values to the right for out of range bins
	for(int iy=5; iy<=yaxis->GetNbins()-5; iy++){
	
		// The right-hand side of the plot will naturally have zeros due to 
		// geometrical acceptance. We don't want to bleed into that region
		// so we need to find the limit where to stop bleeding.
		int xbin_max;
		for(xbin_max=xaxis->GetNbins(); xbin_max>2; xbin_max--){
			double z = h->GetBinContent(xbin_max,iy);
			if(z>=zmin && z<=zmax)break;
		}

		for(int ix=5; ix<xbin_max-5; ix++){
			double z = h->GetBinContent(ix,iy);
			double z1 = h->GetBinContent(ix+2,iy);
			double z2 = h->GetBinContent(ix-2,iy);
			double z3 = h->GetBinContent(ix,iy+2);
			double z4 = h->GetBinContent(ix,iy-2);
			double zavg = (z1+z2+z3+z4)/4.0;
			double zrms = sqrt((pow(z1-zavg,2.0) + pow(z2-zavg,2.0) + pow(z3-zavg,2.0) + pow(z4-zavg,2.0))/4.0);

			if(zavg==0.0)continue;
			if((zrms/zavg)>0.2)continue; // make sure zavg seems to be made from values in reasonable agreement

			if(fabs((z-zavg)/zavg)>0.5){
				h->SetBinContent(ix,iy, zavg);
				Nbins_touched++;
			}
		}
	}
	
	double Nbins_total = (double)xaxis->GetNbins()*(double)yaxis->GetNbins();
	double f = (double)Nbins_touched/Nbins_total;
	cout<<"Number of bins touched for "<<h->GetName()<<": "<<Nbins_touched<<" ("<<f*100.0<<"%)"<<endl;
}