#ifndef SHOWERPDFINCAA
#define SHOWERPDFINCAA

#include "foreach.hh"
#include "rootutil.hh"
#include "util.hh"

#include "TH3D.h"
#include "TFile.h"
#include "Math/IFunction.h"


struct ShowerPdfBase
{
	string name;
	int verb;

	virtual string desc() { return "ShowerPdfBase"; }
	virtual void init()   { cout << "ShowerPdfBase::init" << endl; }
	virtual void init_event( Det& det,
	                         vector<Hit>& hits,
	                         Vec vertex     = Vec(),
	                         double maxdist = 0 )
	{
		cout << "ShowerPdfBase::initev" << endl;
	}

	virtual double eval (Trk& t ) {return 0;}
	virtual ~ShowerPdfBase() {}

	ShowerPdfBase() : verb(0) {}
};


//----------------------------------------------------
// M-estimator PDF
//----------------------------------------------------

struct MestShowerPdf : public ShowerPdfBase
{
	vector<Hit>* hits_;

	double constant;
	double power;

	string desc() { return "m-estimate";}

	MestShowerPdf()
	{
		constant = 1;
		power    = 0.5;
		name     = "MestShowerPdf";
	}

	void init()
	{
		cout << " this is " << name << endl;
	}

	void init_event( Det& det,
	                 vector<Hit>& hits,
	                 Vec vertex     = Vec(),
	                 double maxdist = 0
	               )
	{
		if ( maxdist > 0 ) fatal ("error: MestShowerPdf cannot be called with maxdist>0");

		hits_ = &hits;
	}

	double eval( Trk& trk )
	{
		if (verb > 1) dbg( str(trk) ) ;

		double m = 0;

		foreach ( hit, *hits_ )
		{
			const double r  = (hit.pos - trk.pos).len();
			const double t = hit.t - trk.t - r / v_light;
			m += hit.a * pow( constant + t * t , power );
		}
		return m;
	}

	ClassDef( MestShowerPdf, 1 )
};


//----------------------------------------------------
// Digital Poisson pdf based on P(hit/no hit)
//----------------------------------------------------

/*! A PDF representing the probablity to have either zero hits or more than one hit */

struct DigitalShowerPdf : public ShowerPdfBase
{
	double k40prob;
	double minlogp, maxlogp;
	double minp, maxp;

	bool dbg_probs;
	bool do_clamp;
	bool force_precompute;

	// histograms for the PDF
	string infile;
	TH3D H3;
	TH3D G3;

	// event data
	vector<DomSum> domsums;
	vector<Fdom> fdoms; // pre-computed per-dom information for fast eval
	TH1D hz0;


	enum Eval_mode { plain, quick , hybrid , diagnostic};
	Eval_mode eval_mode;

	string desc() { return "pdf describing P(hit/nohit)";}

	DigitalShowerPdf(const string infile_str = "$AADIR/rec/pdf_data/genhen_JUNE_2___.hist.root") : dbg_probs(false)
	{
		name   = "DigitalShowerPdf";
		infile = infile_str;
		set_minmax_prob( 0.0001 );

		do_clamp = false;
		k40prob = 0.005; // for 5 khz rate and a 1microsecond window
		eval_mode = quick;
	}

	void set_minmax_prob( double pmin )
	{
		minlogp = log( pmin );
		maxlogp = log( 1 - pmin );
		minp = exp( minlogp );
		maxp = exp( maxlogp );
	}

	void init()
	{
		TFile f( infile.c_str() );

		getobj( f, "h3xa" ,    H3 );
		getobj( f, "h3normx" , G3 );

		H3.Divide(&G3);

		if (false ) // for when H3 is not mu, but 1-P0
			for (int ix = 1; ix < H3.GetNbinsX() + 1; ix++)
				for (int iy = 1; iy < H3.GetNbinsY() + 1; iy++)
					for (int iz = 1; iz < H3.GetNbinsZ() + 1; iz++)
					{
						int b = H3.GetBin(ix, iy, iz);
						double p0 = 1 - H3.GetBinContent( b );
						double pp = p0;
						double eps = 1e-30;
						if (p0 < eps )     p0 = eps;
						if (p0 > 1 - eps)    p0 = 1 - eps;
						double mu = -log( p0 );
						if (mu < 0 ) fatal ("highscool math failure.");
						if ( verb > 0 ) print (" transform to mu", ix, iy, iz, "   ",  pp, p0, mu );
						H3.SetBinContent( b, mu );
					}
	}

	void dump()
	{
		cout << "digipdf at" << this << endl;
		cout << " domsums " << domsums.size() << "   " << &domsums << endl;
		cout << " fdoms   " << fdoms.size() << "   " << &domsums << endl;
	}


	/*! Function to be called for every new event. Sets the domsums. */

	void init_event( Det& det,
	                 vector<Hit>& hits,
	                 Vec vertex     = Vec(),
	                 double maxdist = 0 )
	{
		domsums.clear();
		foreach_map( id , pmtptr, det.pmts )
		{
			(void) id;
			pmtptr->flag = 0;
		}

		foreach ( h, hits )
		{
			if ( det.pmts.find (  h.pmt_id ) == det.pmts.end() )
			{
				fatal("cannot find pmt with id", h.pmt_id );
			}
			det.pmts[ h.pmt_id ] -> flag = 1;
		}

		foreach_map( domid, dom , det.doms )
		{
			(void) domid;

			if (maxdist > 0 )
			{
				double dist = (vertex - dom.pos).len();
				if ( dist > maxdist ) continue;
			}

			DomSum D;
			D.pos = dom.pos;

			foreach ( pmt, dom.pmts )
			{
				if (pmt.flag) D.hit_pmts.push_back( pmt.dir );
				else          D.unhit_pmts.push_back( pmt.dir );
			}

			domsums.push_back(D);
		}
		// print ("maxdist=",maxdist, "selected",domsums.size(),"doms out of " , det.doms.size() , " pos=", vertex );

		precompute( vertex );
	}


	double get_h3_value( TH3D& H, double x, double y, double z )
	{
		const int interpolation_policy = 1;

		switch ( interpolation_policy )
		{
		case 0 :  // no interpolation, just bincontent
		{
			const int ix = H.GetXaxis() -> FindBin ( x );
			const int iy = H.GetYaxis() -> FindBin ( y );
			const int iz = H.GetZaxis() -> FindBin ( z );
			return H.GetBinContent( H.GetBin(ix, iy, iz) );
		}

		case 1: // trilinear interploation (by ROOT)
			clamp_bincenter( x, H.GetXaxis() );
			clamp_bincenter( y, H.GetYaxis() );
			clamp_bincenter( z, H.GetZaxis() );
			return H.Interpolate( x, y, z );

		case 2:
			print ("spline interpolation not available in this verion");
		//return interpolate_TH3( H, x,y,z );

		default:
			fatal("wrong interpolation_policy", interpolation_policy);
		}
	}

	double eval( Trk& trk )
	{
		switch ( eval_mode )
		{
		case plain      : return eval_plain( trk );
		case quick      : return eval_quick( trk );
		case hybrid     : return eval_hybrid( trk );
		case diagnostic : return eval_diagnostic( trk );
		}
		fatal("unknown eval_mode");
	}


	double eval_diagnostic( Trk& trk )
	{
		double Lp = eval_plain( trk );
		double Lf = eval_quick( trk );
		double Lh = eval_hybrid( trk );
		print ("likelihoods ", Lp, Lf, Lh );
		return Lf;
	}

	double eval_plain( Trk& trk )
	{
		//TIMER()

		if (trk.E > 1e14 ) trk.E = 1e14;

		const double alpha = trk.E * 1e-6; // energy scale factor ( H3 is for 1 PeV ).

		double L = 0;
		double L1(0), L0(0);
		int nhit = 0;

		foreach ( dom, domsums )
		{
			const Vec v = dom.pos - trk.pos;
			const double r = v.len();
			const Vec dirv = v / r;

			const double z = trk.dir.dot( dirv );

			{
				TIMER( hit - pmt - loop )

				foreach ( pmtdir, dom.hit_pmts )
				{
					nhit ++;
					const double a   = dirv.dot ( pmtdir );
					const double mu1 = get_h3_value( H3, r, z, a );
					const double mu  = alpha * mu1 + k40prob;

					double logp1 = log( 1 - exp( - mu ) );

					if (do_clamp) clamp( logp1, minlogp, maxlogp );
					L -= logp1;
					L1 -= logp1;
				}
			}



			{
				TIMER( unhit - pmt - loop )

				foreach ( pmtdir, dom.unhit_pmts )
				{
					const double a   = dirv.dot ( pmtdir );
					const double mu1 = get_h3_value( H3, r, z, a );
					const double mu  = alpha * mu1 + k40prob;

					// note : - L = alpha * sum( mu1 )  + N*k40prob;
					// can define, for all empty pmts on each dom:
					// sum( f(z) ) = F(z)

					double logp0 = -mu;
					if (do_clamp) clamp( logp0, minlogp, maxlogp );
					L -= logp0;
					L0 -= logp0;
				}
			}

			if ( L != L )
			{
				print("likelihood is nan @ r,z,a =", r, z, 0 );
				print( str(trk) );
			}

		} // dom

		return L;
	}

	bool precompute_if_needed( const Vec& pos )
	{
		static Vec olpos;

		if ( pos != olpos || force_precompute )
		{
			precompute( pos );
			olpos = pos;
			force_precompute = false;
			return true;
		}
		return false;
	}


	void precompute( const Vec& pos ) // fill fdoms
	{
		//TIMER();
		if (verb > 0) print ("precompute________________________________");

		TH1::AddDirectory(0);

		// one prototype histogram (for axis clamping)
		hz0 = histogram_slice_y ( &H3, 100, 0.5 );
		hz0.Scale(0);

		// cout << "resize fdoms to " <<  domsums.size() << " from " << fdoms.size() <<  endl;

		fdoms.resize( domsums.size() , Fdom(hz0) );

		vector<double> va(31);

		enumerate( i, dom, domsums )
		{
			Fdom& fdom = fdoms[i];

			const Vec v = dom.pos - pos;
			const double r = v.len();
			fdom.dir = v / r;

			//---------------------------------------------------------
			// for all unhit pmts, we make a histogram of sum(mu) vs z
			// per dom
			//---------------------------------------------------------

			// fill va.
			const int N = dom.unhit_pmts.size();
			va.resize( N );
			for (int ipmt = 0; ipmt < N; ipmt++) va[ipmt] = dom.unhit_pmts[ipmt].dot ( fdom.dir );

			fdom.hist0.Scale(0);
			interpol_xz( &H3, r , va , &fdom.hist0 );
			fdom.spline0 = Spline( fdom.hist0 );

			// now the non-emtpy doms
			fdom.hit_hists.clear();
			fdom.hit_splines.clear();

			foreach ( pmtdir, dom.hit_pmts )
			{
				const double a = pmtdir.dot ( fdom.dir );
				fdom.hit_hists.push_back(  histogram_slice_y ( &H3, r, a ) );
				fdom.hit_splines.push_back( Spline( *fdom.hit_hists.rbegin() ) );
			}

		} // domsums loop
	}


	void eval_gradient_quick( Trk& t,
	                          double& f, double gradient[3] )
	{
		double theta = acos( t.dir.z );
		double phi = atan2( t.dir.y, t.dir.x );
		eval_gradient_quick( theta, phi, log10(t.E) , f , gradient );
	}



	void num_gradient( double theta, double phi, double logE,
	                   double gradient[3])
	{
		double h = 1e-7;

		double f1, f2;
		double df1[3], df2[3];
		double E = logE;

		eval_gradient_quick2( theta + h, phi, E , f1 , df1 );
		eval_gradient_quick2( theta - h, phi, E , f2 , df2 );
		gradient[0] = -(f2 - f1) / (2 * h);

		eval_gradient_quick2( theta, phi + h, E , f1 , df1 );
		eval_gradient_quick2( theta, phi - h, E , f2 , df2 );
		gradient[1] = -(f2 - f1) / (2 * h);

		eval_gradient_quick2( theta, phi, E + h , f1 , df1 );
		eval_gradient_quick2( theta, phi, E - h , f2 , df2 );
		gradient[2] = -(f2 - f1) / (2 * h);
	}

	void fill_hessian( double* hessian, double theta, double phi, double logE , double h = 1e-7 )
	{
		double g1[3];
		double g2[3];
	
		for (int i=0; i<3; i++)
		{
			double d1 = i==0? h : 0 ;
			double d2 = i==1? h : 0 ;
			double d3 = i==2? h : 0 ;
			double f;

			eval_gradient_quick( theta + d1, phi + d2, logE + d3, f, g1 );
			eval_gradient_quick( theta - d1, phi - d2, logE - d3, f, g2 );

			for( int j = 0; j<3; j++) hessian[i*3+j] = (g1[j]-g2[j])/(2*h);
		}
	}


	/*! numerically evaluate hessian (based on analytical gradient) */
	    
	vector<double> eval_hessian ( double theta, double phi, double logE , double h=1e-7)
	{
		vector<double> r(9);
		fill_hessian( &r[0], theta, phi, logE, h );	
		return r;
	}


	void eval_gradient_quick( double theta, double phi, double logE,
	                          double& f, double gradient[3] )
	{

		eval_gradient_quick2( theta, phi, logE, f, gradient );

		if ( verb )
		{
			double gr[3];
			num_gradient( theta, phi, logE , gr );
			print ("eval gradient", theta, phi, logE , " -> ", f ,
			       "\t grad=", gradient[0], gradient[1],
			       "\t num =", gr[0], gr[1] );
		}

	}

	void eval_gradient_quick1( double theta, double phi, double logE,
	                           double& f, double gradient[3] )
	{
		const double energy_scale_factor = 1e-6;

		double E     = pow(10, logE);

		if ( E < 1e-12 )   E = 1e-12;
		if ( E > 1e12  )   E = 1e12;

		const double alpha = energy_scale_factor * E;

		const double sint = sin(theta);
		const double sinf = sin(phi);
		const double cost = cos(theta);
		const double cosf = cos(phi);

		const double zmin = hz0.GetBinCenter( 1 );
		const double zmax = hz0.GetBinCenter( hz0.GetNbinsX() );

		f = gradient[0] = gradient[1] = gradient[2] = 0;

		foreach ( fdom, fdoms )
		{
			// z= fdom.dir (dot) trk.direction

			double z =
			    fdom.dir.x * sint * cosf +
			    fdom.dir.y * sint * sinf +
			    fdom.dir.z * cost;

			double dz_dtheta =
			    fdom.dir.x * cost * cosf +
			    fdom.dir.y * cost * sinf -
			    fdom.dir.z * sint;

			double dz_dphi =
			    - fdom.dir.x * sint * sinf +
			    fdom.dir.y * sint * cosf;

			if ( z < zmin ) { z = zmin ; dz_dtheta = dz_dphi = 0 ; }
			if ( z > zmax ) { z = zmax ; dz_dtheta = dz_dphi = 0 ; }

			double L(0), dL_dz(0), dL_dlogE(0);

			foreach ( h, fdom.hit_hists )
			{
				double f, df_dz;
				interpolate_derivative( &h, z , f, df_dz);


				const double mu        = k40prob + alpha * f;
				const double dmu_dz    = alpha * df_dz;
				const double dmu_dE    = f * energy_scale_factor;
				const double dmu_dlogE = dmu_dE * ( log(10) * E );

				const double exp_mu = exp(-mu);

				const double p        = 1.0 - exp_mu;
				const double dp_dz    = exp_mu * dmu_dz;
				const double dp_dlogE = exp_mu * dmu_dlogE;

				L     -= log(p) ;
				dL_dz -= dp_dz / p;
				dL_dlogE -= dp_dlogE / p;

				if (L != L)
				{
					print ("-----------------------------");
					print ( alpha, df_dz, f, mu, exp_mu, p );

					dbg( L );
					fatal("nan");
				}
			}



			// unhit pmts
			const int nempty = 31 - fdom.hit_hists.size();
			if (nempty > 0)
			{
				double f, df_dz;
				interpolate_derivative( &fdom.hist0, z , f, df_dz);

				const double p     = k40prob * nempty + alpha * f;
				const double dp_dz = alpha * df_dz;
				const double dp_dlogE = f * energy_scale_factor * E * log(10);

				L     += p;
				dL_dz += dp_dz;
				dL_dlogE += dp_dlogE;
			}

			f += L;
			gradient[0] += dL_dz * dz_dtheta;
			gradient[1] += dL_dz * dz_dphi;
			gradient[2] += dL_dlogE;
		}



	}



	/*! Version that uses spline */

	void eval_gradient_quick2( double theta, double phi, double logE,
	                           double& f, double gradient[3] )
	{
		//TIMER();

		const double energy_scale_factor = 1e-6;

		double E     = pow(10, logE);

		if ( E < 1e-12 )   E = 1e-12;
		if ( E > 1e12  )   E = 1e12;

		const double alpha = energy_scale_factor * E;

		const double sint = sin(theta);
		const double sinf = sin(phi);
		const double cost = cos(theta);
		const double cosf = cos(phi);

		const double zmin = hz0.GetBinCenter( 1 );
		const double zmax = hz0.GetBinCenter( hz0.GetNbinsX() );

		f = gradient[0] = gradient[1] = gradient[2] = 0;

		foreach ( fdom, fdoms )
		{
			// z= fdom.dir (dot) trk.direction

			double z =
			    fdom.dir.x * sint * cosf +
			    fdom.dir.y * sint * sinf +
			    fdom.dir.z * cost;

			double dz_dtheta =
			    fdom.dir.x * cost * cosf +
			    fdom.dir.y * cost * sinf -
			    fdom.dir.z * sint;

			double dz_dphi =
			    - fdom.dir.x * sint * sinf +
			    fdom.dir.y * sint * cosf;

			if ( z < zmin ) { z = zmin ; dz_dtheta = dz_dphi = 0 ; }
			if ( z > zmax ) { z = zmax ; dz_dtheta = dz_dphi = 0 ; }

			double L(0), dL_dz(0), dL_dlogE(0);

			foreach ( h, fdom.hit_splines )
			{
				double f, df_dz;

				h.eval_derivative( z, f, df_dz );
				//interpolate_derivative( &h, z , f, df_dz);

				//print ("spline eval ", z, f, df_dz );

				const double mu        = k40prob + alpha * f;
				const double dmu_dz    = alpha * df_dz;
				const double dmu_dE    = f * energy_scale_factor;
				const double dmu_dlogE = dmu_dE * ( log(10) * E );

				const double exp_mu = exp(-mu);

				const double p        = 1.0 - exp_mu;
				const double dp_dz    = exp_mu * dmu_dz;
				const double dp_dlogE = exp_mu * dmu_dlogE;

				L     -= log(p) ;
				dL_dz -= dp_dz / p;
				dL_dlogE -= dp_dlogE / p;

				if (L != L)
				{
					print ("-----------------------------");
					print ( alpha, df_dz, f, mu, exp_mu, p );

					dbg( L );
					fatal("nan");
				}
			}



			// unhit pmts
			const int nempty = 31 - fdom.hit_hists.size();
			if (nempty > 0)
			{
				double f, df_dz;

				fdom.spline0.eval_derivative( z, f, df_dz );

				const double p     = k40prob * nempty + alpha * f;
				const double dp_dz = alpha * df_dz;
				const double dp_dlogE = f * energy_scale_factor * E * log(10);

				L     += p;
				dL_dz += dp_dz;
				dL_dlogE += dp_dlogE;
			}

			f += L;
			gradient[0] += dL_dz * dz_dtheta;
			gradient[1] += dL_dz * dz_dphi;
			gradient[2] += dL_dlogE;
		}
	}

	double eval_quick( Trk& trk )
	{
		precompute_if_needed( trk.pos );
		//TIMER();

		double L = 0;

		if (trk.E > 1e14 ) trk.E = 1e14;

		const double alpha = trk.E * 1e-6; // energy scale factor ( H3 is for 1 PeV ).

		const int nterms      = 1; // code is not correct for values > 1
		int termcounter = 0;
		double P = 1;
		int nhit = 0;

		foreach ( fdom, fdoms )
		{
			double z = fdom.dir.dot( trk.dir );
			clamp_bincenter( z, hz0.GetXaxis());

			foreach ( h, fdom.hit_hists )
			{
				nhit++;
				const double mu = k40prob + alpha * interpolate( &h, z );

				double p = 1.0 - exp( -mu );
				if (do_clamp) clamp( p, minp, maxp );
				P *= p;


				if ( ++termcounter == nterms )
				{
					L  -= log(P);
					P = 1.0;
					termcounter = 0;
				}
			}

			// unhit pmts
			const int nempty = 31 - fdom.hit_hists.size();
			if (nempty)
			{
				const double p = k40prob * nempty + alpha * interpolate( &fdom.hist0, z );
				L += p;
			}

		}// fdomloop

		return L;
	}


	double eval_hybrid( Trk& trk )
	{
		precompute_if_needed(trk.pos);

		//TIMER();

		double L = 0;
		double L0 = 0;
		double L1 = 0;
		if (trk.E > 1e14 ) trk.E = 1e14;

		const double alpha = trk.E * 1e-6; // energy scale factor ( H3 is for 1 PeV ).

		int nhit = 0;

		{
			TIMER( Hit - pmts );

			foreach ( dom, domsums )
			{
				const Vec v    = dom.pos - trk.pos;
				const double r = v.len();
				const Vec dirv = v / r;
				const double z = trk.dir.dot( dirv );

				foreach ( pmtdir, dom.hit_pmts )
				{
					nhit ++;
					const double a   = dirv.dot ( pmtdir );
					const double mu1 = get_h3_value( H3, r, z, a );
					const double mu  = alpha * mu1 + k40prob;

					double logp1 = log( 1 - exp( - mu ) );

					if (do_clamp) clamp( logp1, minlogp, maxlogp );
					L -= logp1;
					L1 -= logp1;
				}
			}

		} // timer


		{
			TIMER (empty - pmts )

			foreach ( fdom, fdoms )
			{
				const double z = trk.dir.dot( fdom.dir );
				const int nempty = 31 - fdom.hit_hists.size();
				double p = k40prob * nempty + alpha * interpolate( &fdom.hist0, z );
				L += p;
				L0 += p;
			}
		} // timer

		return L;
	}


	ClassDef( DigitalShowerPdf, 1);

};


#include "Math/IFunction.h"

struct GradFunc : public ROOT::Math::IGradientFunctionMultiDim
{
	DigitalShowerPdf* pdf;

	virtual IBaseFunctionMultiDimTempl<double> *Clone() const {return 0;}

	virtual unsigned int 	NDim () const { return 3; }

	virtual double DoEval(const double *x) const
	{

		double f;
		double df[3];
		pdf -> eval_gradient_quick ( x[0], x[1], x[2], f, df );

		print ("do_eval", f);
		return f;
	}

	virtual double DoDerivative(const double *x,
	                            unsigned int icoord) const
	{

		double f;
		double df[3];
		pdf -> eval_gradient_quick ( x[0], x[1], x[2], f, df );

		print ( "do_derivative", f, df[0], df[1], df[2] );
		return df[icoord];

	}

	virtual void FdF (const double* x, double& f, double* df) const
	{
		pdf -> eval_gradient_quick ( x[0], x[1], x[2], f, df );

		print ( "fdf", f, df[0], df[1], df[2] );
	}


};

struct Likresult
{
	double L; // total likelihood
	double L1; // likelihood of hit pmts
	double L0; // likelihood of unhit pmts
};

#endif