#include "TCanvas.h"
#include "TH3D.h"
#include "TSpline.h"
#include "TFitResult.h"
#include "TRandom2.h"
#include <algorithm>
#include "stringutil.hh"

/* new rules:
   a    = angle in degrees between nu and rec.
   loga = log10( a )
*/


inline void check( TH1& h , string ss ) {
	if ( h.GetNbinsX() < 1) print ( " ---- histogram is dummy ---- ", ss );
	for (int b = 1 ; b < h.GetNbinsX(); b++ ) {
		double c = h.GetBinContent( b );
		if ( ::isnan(c) ) {
			print ( " ---- histogram has NaN ---- ", ss );
			print ( h.GetName() );
			for (int bb = 1 ; bb < h.GetNbinsX(); bb++ ) {
				cout << h.GetBinContent(bb) << " ";
			}
			fatal("---------------------------------------");
		
		}	
	}
}


/*! The Point spread function, as a function of the true energy.
    The point spread is a PDF; as such it's integral is 1. (it does not
    take into account the number of events -- that is the job of Aeff)
 */
struct PsfParam
{
	string name;
	vector<TSpline3> splines ;           // of dP/dOmega
	vector<TH1D>     slices_loga;        // of dP/dloga, smoothed
	vector<TH1D>     slices_loga_org;    // of dP/dloga, not smoothed
	vector<TH1D>     slices_omega;       // of dP/dOmega   , smoothed
	vector<TH1D>     slices_omega_org;   // untreated version

	int NbinslogE;
	int minbin, maxbin;
	double logE_min;
	double logE_max;

	/*! Jacobian d Omega / d log(a), with a angle in degrees  */
	double d_omega_d_loga ( double loga )
	{
		double a = pow( 10, loga) * pi / 180;
		return sin(a)  * 2 * pi * log(10) * a;
	}

	/* Jacobian d log(a) / d Omega , with a angle in degrees */
	double d_loga_d_omega ( double loga )
	{
		return 1.0 / d_omega_d_loga ( loga );
	}

	/* Integrate the PSF over the sphere at log E */
	double integrate_sphere( double logE , double max_angle = pi )
	{
		double dla = 0.01;
		double c = 0;

		for ( double la = -3; la < log10(max_angle*180/pi) ; la += dla )
		{
			double P = eval( logE, la );
			double y = d_omega_d_loga ( la ) * P * dla;
			c += y;
		}
		return c;
	}

	/* Integrate the PSF over the sphere between min_angle and max_angle at log E */
	double integrate_sphere_angle_range( double logE , double min_angle, double max_angle = pi )
	{
		double dla = 0.001;                          // step size
	
		double mina = log10( 1e-3*pi/180 );
		double maxa = log10( 180/pi );
		double log_min_angle = log10( min_angle *180/pi );
		double log_max_angle = log10( max_angle *180/pi );

		double c = 0;
		for ( double la = mina ; la < maxa ; la += dla )
		{
			// the loop is done this way so that the same points are always evaluated
			// so that sub-ranges in angle correctly add up.

			if ( la <= log_min_angle || la > log_max_angle ) continue;

		// if (min_angle == 0) min_angle = 1e-3*pi/180;

		// for ( double la = log10(min_angle*180/pi); la < log10(max_angle*180/pi) ; la += dla )
		// {
			double P = eval( logE, la );
			double y = d_omega_d_loga ( la ) * P * dla;
			c += y;
			// cout << "c " << c << " la " << la << " P " << P << " y " << y << endl;
		}
		return c;
	}

	/* return the probability the neutrino will be < angle_max */
	double fraction_below_angle( double logE, double angle_max )
	{
		double c = integrate_sphere( logE );
		//print ("total psf integral", c );
		if ( c == 0 ) return 0;
		else return integrate_sphere( logE, angle_max ) / c;
	}

	/* return the probability the neutrino will be between angle_min and angle_max */
	double fraction_between_angle( double logE, double angle_min, double angle_max )
	{
		double c = integrate_sphere( logE );
		//print ("total psf integral", c );
		if ( c == 0 ) return 0;
		else return integrate_sphere_angle_range( logE, angle_min, angle_max ) / c;
	}

	double integrate_sphere( TH1D& dP_dOmega )
	{
		double c;
		for ( int b = 1; b <= dP_dOmega.GetNbinsX() ; b++ )
		{
			double loga = dP_dOmega.GetBinCenter(b);
			c += d_omega_d_loga ( loga ) * dP_dOmega.GetBinContent(b) * dP_dOmega.GetBinWidth(b);
		}
		return c;
	}


	void test()
	{
		for (double le = 2; le < 8; le += 0.33 )
		{
			print ( le, integrate_sphere( le ));
		}
	}

	double dP_dOmega(  double logE_true, double loga, double cos_theta_true = -99 )
	{
		return eval( logE_true, loga );
	}

	double dP_dloga( double logE_true, double loga, double cos_theta_true = -99)
	{
		return eval ( logE_true, loga ) / d_loga_d_omega( loga );
	}

	TH1D get_hist_dP_dOmega( double logE_true )
	{
		TH1D r("r", "r", 100, -3, 3);
		for (int b = 1; b <= r.GetNbinsX() ; b++ )
		{
			double log_a = r.GetBinCenter(b);
			r.SetBinContent( b, dP_dOmega( logE_true, log_a ));
		}
		return r;
	}

	TH1D get_hist_dP_dloga( double logE_true )
	{
		TH1D r("r", "r", 100, -3, 3);
		for (int b = 1; b <= r.GetNbinsX() ; b++ )
		{
			double log_a = r.GetBinCenter(b);
			r.SetBinContent( b, dP_dloga( logE_true, log_a ));
		}
		return r;
	}

	double eval( double logE, double loga )
	{
		int b0, b1; double f0, f1;
		get_interpolation_consts( NbinslogE, logE_min, logE_max, logE,
		                          b0, b1, f0, f1 );
		clamp( b0, minbin, maxbin );
		clamp( b1, minbin, maxbin );
		if ( loga < -3 )        loga = -3;
		if ( loga > log10(180)) loga = log10(180);
		double r = f0 * splines[b0].Eval( loga ) + f1 * splines[b1].Eval( loga );
		if (::isnan(r)) {
			fatal("nan in PsfParam::eval", b0,b1, f0, f1 );
		}
		return r;
	}

	double get_random_loga( double logE )
	{
		int b0, b1; double f0, f1;
		get_interpolation_consts( NbinslogE, logE_min, logE_max, logE,
		                          b0, b1, f0, f1 );
		if ( b0 < minbin ) b0 = minbin;
		//if ( b0 > maxbin ) b0 = maxbin;
		//if ( b1 < minbin ) b1 = minbin;
		if ( b1 > maxbin ) b1 = maxbin;
		const int b = gRandom->Rndm() < f0 ? b0 : b1;
		return slices_loga[ b ].GetRandom();
	}

	void get_random(double logE, double dec_in, double ra_in, double& dec_out, double& ra_out )
	{
		const double a = pow(10, get_random_loga( logE )) * pi / 180;
		const double theta_in = pi / 2 - dec_in;
		Vec v; v.set_angles( theta_in, ra_in );
		Mat M; M.set(v);
		Vec r; r.set_angles( a, gRandom->Rndm() * 2 * pi );
		Vec rr = M * r;
		dec_out = pi / 2 - rr.theta();
		ra_out  = rr.phi();
	}

	/*! average content of bins b0 to b1 (inclusive) */
	double average( TH1D h, int b0, int b1 )
	{
		if (b0 > b1) print("error: averaging over wrong bin set",b0,b1);
		check( h , " average ");
		double s = 0;
		int n = 0;
		double w_s = 0;
		for ( int b = b0 ; b <= b1; b++)
		{
			double y    = h.GetBinContent(b);
			double loga = h.GetBinCenter(b);
			double a    = pow(10, loga);
			double weight = 2 * pi * (1 - cos(a)); //omega
			//cout<<"in PsfParam.hh nbin="<< h.GetBin(b)<<" loga="<<loga<<" , a="<<a<<" , weight(omega)="<<weight<<" ,y="<<y<<endl;
			if ( y == 0 ) continue;
			n++;
			//s += y;
			s += y * weight;
			w_s += weight;
		}
		if (w_s == 0) return 0;  //all y were 0 so average is 0 too
		//return s / n  ;
		return s / w_s ;
	}

	/* return true if the histogram is decreasing between b0 and FindBin(X) */
	bool is_good( TH1D& h , int b0 , double X )
	{
		double a = average( h, 1, b0 );
		if ( a < h.GetBinContent( b0 + 1 ) ) return false;
		int b1 = h.FindBin( X );
		for ( int b = b0 + 1 ; b <= b1 ; b++ )
		{
			//print ( "-->", b, h->GetBinContent(b) );
			double c = h.GetBinContent( b );
			double c1 = h.GetBinContent( b + 1);
			if ( c < c1 )
			{
				//print ("bin", b+1,"is larger than ", b , c, c1, h->GetBinCenter(b) );
				return false;
			}
		}
		return true;
	}

	/*! set bins b0 to b1 (inclusive) to their average */
	void set_to_average( TH1D& h , int b0, int b1 )
	{
		double a = average( h, b0, b1 );
		
		//for ( int b = b0 ; b<=b1; b++)
		for ( int b = 1 ; b <= b1; b++)
		{
			h.SetBinContent(b, a );
		}
	}

	void source_extension_smearing(TH2D &h, double width)//width in deg
	{
		TRandom2 rnd;
		rnd.SetSeed(0);
		int nbinx = h.GetXaxis()->GetNbins();
		TH2D h_tmp = h;  h_tmp.Reset();
		int ngen = 10000;
		for (int i = 1; i <= nbinx; i++) {
			string name = "projy_" + to_string(i);
			double w = h.ProjectionY(name.c_str(), i, i)->Integral();
			if (h.ProjectionY(name.c_str(), i, i)->GetEntries() != 0) {
				for (int nnn = 0; nnn < ngen; nnn++) {
					double loga = h.ProjectionY(name.c_str(), i, i)->GetRandom();
					double a = pow(10, loga) / deg; //rad
					double Psi = width / deg * sqrt( -2 * log(rnd.Rndm()) ); //rad
					Vec Vrec; Vrec.set_angles( a, rnd.Rndm() * 2 * pi );
					Vec Vsrc; Vsrc.set_angles( Psi, rnd.Rndm() * 2 * pi );
					double Csi = angle_between(Vrec, Vsrc); //rad
					h_tmp.Fill(h.GetXaxis()->GetBinCenter(i), log10(Csi * deg), w / (double)ngen);
				}
			} else {
				double loga = -7;
				h_tmp.Fill(h.GetXaxis()->GetBinCenter(i), loga);
			}
		}
		h = h_tmp;
	}

	void source_extension_smearing_analytic(TH2D &h, double width)//width in deg
	{
		TRandom2 rnd;  rnd.SetSeed(0);
		int nbinx = h.GetXaxis()->GetNbins();
		int nbinloga = h.GetYaxis()->GetNbins();
		double xmin = h.GetYaxis()->GetXmin();
		double xmax = h.GetYaxis()->GetXmax();
		TH2D h_tmp = h;  h_tmp.Reset();
		double atot = h.Integral();
		//time_t timer1 = time(NULL);
		int npoints = 100;
		double tofill, tot;
		double fdPdlogaval, dloga, dlogPsi, dPhidCsi;
		double loga, logPsi, logCsi;
		double a, Psi, Csi, minlogPsi, maxlogPsi, cosPhi;
		// gaussian shape
		TF1 * fdPdlogPsi = new TF1("hdPdlogpsi", [&](double * x, double * par) {
			double xxx = pow(10, x[0]);
			double_t arg = (xxx - par[1]) / par[2];
			double res = arg * arg * exp(-0.5 * arg * arg);
			return par[0] * res;
		}, xmin, xmax, 3);
		fdPdlogPsi->SetParameter(0, 0.135); // to check
		fdPdlogPsi->SetParameter(1, 0.);
		fdPdlogPsi->SetParameter(2, width);
		for (int i = 1; i <= nbinx; i++) {
			string name = "projy_" + to_string(i);
			//double w=h.ProjectionY(name.c_str(),i,i)->Integral();
			if (h.ProjectionY(name.c_str(), i, i)->GetEntries() != 0) {
				//   **********
				for (int ilogCsi = 1; ilogCsi <= nbinloga; ilogCsi++) {
					logCsi = h.ProjectionY(name.c_str(), i, i)->GetXaxis()->GetBinCenter(ilogCsi);
					Csi = pow(10, logCsi);
					tot = 0;
					dloga = (xmax - xmin) / static_cast<double>(npoints);
					for (int iloga = 1; iloga < nbinloga; iloga++) {
						loga = h.ProjectionY(name.c_str(), i, i)->GetXaxis()->GetBinCenter(iloga);
						a = pow(10, loga); //deg
						fdPdlogaval = h.ProjectionY(name.c_str(), i, i)->GetBinContent(iloga);
						if (a - Csi != 0) minlogPsi = log10(fabs(a - Csi));
						else minlogPsi = xmin;
						if (a + Csi < 180) maxlogPsi = log10(a + Csi);
						else maxlogPsi = xmax;
						dlogPsi = (maxlogPsi - minlogPsi) / static_cast<double>(npoints);
						for (int ipsi = 0; ipsi < npoints; ipsi++) {
							logPsi = minlogPsi + dlogPsi * (ipsi + 0.5);
							Psi = pow(10, logPsi); //deg
							cosPhi = (cos(Csi / deg) - cos(a / deg) * cos(Psi / deg)) / sin(a / deg) / sin(Psi / deg);
							if (fabs(cosPhi) < 1) {
								dPhidCsi = 1. / sqrt(1 - cosPhi * cosPhi) * sin(Csi / deg) / sin(a / deg) / sin(Psi / deg);
								tofill = fdPdlogaval * fdPdlogPsi->Eval(logPsi) * dPhidCsi * Csi * dloga * dlogPsi;
								if (tofill > 0) tot += tofill;
							}
						}
					}
					h_tmp.AddBinContent(h.GetBin(i, ilogCsi), tot);
				}
				// *****************
			} else {
				double loga = -7;
				h_tmp.Fill(h.GetXaxis()->GetBinCenter(i), loga);
			}
		}
		h = h_tmp;
		h.Scale(atot / h.Integral());
	}




	void source_extension_smearing_disk(TH2D &h, double width)//width in deg
	{

		TRandom2 rnd;
		rnd.SetSeed(0);
		int nbinx = h.GetXaxis()->GetNbins();
		TH2D h_tmp = h;  h_tmp.Reset();
		int ngen = 10000;

		for (int i = 1; i <= nbinx; i++) {
			string name = "projy_" + to_string(i);
			double w = h.ProjectionY(name.c_str(), i, i)->Integral();

			if (h.ProjectionY(name.c_str(), i, i)->GetEntries() != 0) {
				for (int nnn = 0; nnn < ngen; nnn++) {
					double loga = h.ProjectionY(name.c_str(), i, i)->GetRandom();
					double a = pow(10, loga) / deg; //rad
					//double Psi = width / deg * sqrt( -2 * log(rnd.Rndm()) ); //rad
                                	double Psi = width / deg * sqrt(rnd.Rndm() ); //rad
					Vec Vrec; Vrec.set_angles( a, rnd.Rndm() * 2 * pi );
					Vec Vsrc; Vsrc.set_angles( Psi, rnd.Rndm() * 2 * pi );
					double Csi = angle_between(Vrec, Vsrc); //rad
					h_tmp.Fill(h.GetXaxis()->GetBinCenter(i), log10(Csi * deg), w / (double)ngen);
				}
			} else {
				double loga = -7;
				h_tmp.Fill(h.GetXaxis()->GetBinCenter(i), loga);

			}
		}
		h = h_tmp;





	}






  void init( TH3D& h_ANGERR, int mergebins, double src_width = 0 ,bool gaus=true )	//src_width in deg
	{
		// Combine all zenith angles to increase statistics.
		// that is: we assume the PSF does not depend on zenith.
		// Also : rebin energy by factor mergebin.
		TH2D& h_ang_2d = *((TH2D*) h_ANGERR.Project3D("zx"));
		h_ang_2d.SetDirectory(0);
		if (src_width != 0) {
		  if(gaus==true){
		    source_extension_smearing(h_ang_2d, src_width);
		  } else if(gaus==false){
		    source_extension_smearing_disk(h_ang_2d, src_width);		    
		  }else{
		  fatal("Invalid source shape!");

		  }
		}
		TAxis& energy_axis = *h_ANGERR.GetXaxis();
		NbinslogE = energy_axis.GetNbins() / mergebins;
		logE_min  = energy_axis.GetXmin();
		logE_max  = energy_axis.GetXmax();
		// For each energy-bin, we will have a set of 1d histograms
		// that we use to store smooth versions of dP/dloga and dP/dOmega.
		// We also keep the orininal slices.
		slices_loga.resize( NbinslogE );
		slices_loga_org.resize( NbinslogE );
		slices_omega.resize( NbinslogE );
		slices_omega_org.resize( NbinslogE );
		splines.resize( NbinslogE );
		int bb = 1;
		for (int b = 0; b < NbinslogE; b++ )
		{
			h_ang_2d.GetXaxis()->SetRange( bb, bb + mergebins - 1 );
			TH1D h = *h_ang_2d.ProjectionY();
			check( h , " after projection ");
			h.SetDirectory(0);
		
			std::cout << "." << std::flush;
			if ( h.Integral() == 0 ) {
				minbin = b + 1;
			}
			else {
				maxbin = b;
				h.Scale( 1.0 / h.Integral() );
				// at this point, h contains the distribution of loga
				// (log10(angle-nu-rec)).
				// now transform it to dP/dOemga
				TH1D htrans = transform( h );
				check ( htrans, " htrans ");
				slices_loga_org[b]  = h;
				slices_omega_org[b] = to_dP_dOmega(h, true);
				slices_omega[b]     = htrans;
				splines[b]          = to_spline( htrans );
				slices_loga[b]      = to_dP_dloga( htrans ); // for consistency check with input
			}
			bb += mergebins;
		}
	}
	/* Massage a histogram of dP/dloga and attempt to produde a
	   decent dP/dOmega as output.
	   This is a series of add-hoc operations that try to ensure
	   that the function is constant for very small values and
	   alwasy degreasing otherwise.*/
	TH1D transform( TH1D h )
	{
		auto fr = h.Fit("gaus", "SQ0", "", -2, 2.5).Get();
		
		if (fr == nullptr )
		{
			return TH1D();
		}
		double mean  = fr->GetParams()[1];
		double sigma = fr->GetParams()[2];
		
		if ( mean < -3 || mean > 2 ) { // try to rescue very low stat. cases
			mean  = 0.5;
			sigma = 1;
		}
		double x0 = std::max( mean - 2.0 * sigma, -2.0 );
		double x1 = std::min( mean + 2.0 * sigma,  1.0 );
		int b0 = h.FindBin( x0 );
		int b1 = h.FindBin( x1 );	
		while ( h.GetBinContent( b0 ) == 0 && b0 < h.GetNbinsX() ) b0++;
		while ( h.GetBinContent( b1 ) == 0 && b1 >0              ) b1--;
		TH1D h2 = to_dP_dOmega(h);
		h2.GetListOfFunctions()->Clear();

		check(h2, "to_dP_dOmega prior to massaging");
		// print ( " fit : ", mean, sigma, x0, x1 );
		// print ( b0 , h2.FindBin(x0) );
		// print ( b1 , h2.FindBin(x1) );
		
		// find b1 so that, for bin > b1, the histogram is dereasing.
		
		
		for ( ; b0 < b1 ; b0 ++ )
		{
			if (is_good( h2, b0, x1 )) break;
		}
		if ( b0 == b1 )
		{
			print ("failed to find good start bin for ", name );
		}
		// for (int i=1; i < h.GetNbinsX() ;i++ ) {
		// 	print (i, h.GetBinCenter(i), h.GetBinContent(i) );
		// }
		// find b_start, starting bin for the plateau value calculation
		int b_start;
		for ( b_start = 1; b_start < b0 ; b_start ++ )
		{
			if (h2.GetBinContent(b_start) != 0 && h2.GetBinContent(b_start + 1) != 0 && h2.GetBinContent(b_start + 2) != 0) break;
		}
		set_to_average( h2, b_start, b0 );
		// check( h2, "after set to average " + str(b_start) + " " + str( b0 ) );
		double s = integrate_sphere( h2 );
		if ( s == 0 ) {
			fatal (" histogram empty in transform ");
		}
		h2.Scale(1.0 / s);
		// check( h2 , "h2 again ");
		return h2;
	}
	/* convert a histogram of dP_dOmega to dP_dloga */
	TH1D to_dP_dloga( TH1D& h_dP_dOmega  )
	{
		TH1D r = h_dP_dOmega;
		for (int b = 1; b < h_dP_dOmega.GetNbinsX() + 1; b ++ )
		{
			double log_a = h_dP_dOmega.GetBinCenter(b);
			double y = h_dP_dOmega.GetBinContent(b) * d_omega_d_loga ( log_a );
			r.SetBinContent( b, y );
		}
		return r;
	}
	/* convert a histogram of dP_dloga to dP_dOmega */
	TH1D to_dP_dOmega( TH1D& h_dP_dloga, bool scale = false  )
	{
		TH1D r = h_dP_dloga;
		for (int b = 1; b < h_dP_dloga.GetNbinsX() + 1; b ++ )
		{
			double log_a = h_dP_dloga.GetBinCenter(b);
			double y = h_dP_dloga.GetBinContent(b) / d_omega_d_loga ( log_a );
			r.SetBinContent( b, y );
		}
		if (scale) {
			double s = integrate_sphere( r );
			r.Scale(1.0 / s);
		}
		return r;
	}
	TSpline3 to_spline( TH1D& h )
	{
		check( h , "to spline");
		vector<double> xs;
		vector<double> ys;
		for (int b = 1; b <= h.GetNbinsX() ; b++ )
		{
			double bc = h.GetBinCenter(b);
			double y  = h.GetBinContent(b);
			if (::isnan(bc) || ::isnan( y )) {
				fatal("nan in to_spline", b, bc, y );
			}
			if ( bc > log10(180) ) break;
			xs.push_back(bc);
			ys.push_back( h.GetBinContent( b));
		}
		xs.push_back( log10(180) );
		ys.push_back( 0 );
		TSpline3 spline( "s", &(xs[0]), &ys[0], xs.size() , "e1b1"  );
		spline.SetLineColor(4);
		spline.SetLineWidth(2);
		return spline;
	}
	~PsfParam() {
		slices_loga.clear();
		slices_loga_org.clear();
		slices_omega.clear();
		slices_omega_org.clear();
	}
	PsfParam() {}
	PsfParam( TH3D& h_ANGERR, int mergebins )
	{
		init( h_ANGERR, mergebins );
	}
};