
#include <algorithm>
#include <limits>
#include <cmath>
#include <vector>
#include <numeric>
#include <map>

#include "confidence.h"

#include "common/assertion.h"

//Brute-force testing
#include <gsl/gsl_cdf.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_histogram.h>

//Root finding
#include <gsl/gsl_errno.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_roots.h>
using std::vector;

//Kahan Summation
//https://codereview.stackexchange.com/questions/56532/kahan-summation
template <class T, class T2>
typename std::iterator_traits<T>::value_type kahansum(T begin, T end, T2 other) {
	typedef typename std::iterator_traits<T>::value_type real;
	real sum = real();
	real running_error = real();

	for ( ; begin != end; ++begin) {
		real difference = *begin - running_error;
		real temp = sum + difference;
		running_error = (temp - sum) - difference;
		sum = temp;
		*other=sum;
		++other;
	}
	return sum;
}

void poissonConfidenceLimit(float observedCounts, float alpha, float &lBound, float &uBound)
{
	ASSERT(observedCounts >=0);
	ASSERT(alpha>0 && alpha < 1);

	float alphaOne=(1.0f-alpha)/2.0;
	float alphaTwo = 1.0-alphaOne;

	if(!observedCounts)
		lBound=0;
	else
		lBound=gsl_cdf_chisq_Pinv(alphaOne,2.0*observedCounts)/2.0f;

	uBound=gsl_cdf_chisq_Pinv(alphaTwo,2.0*(observedCounts+1))/2.0f;    
}


//Brute-force poisson ratio confidence estimator
// Biased, as it enforces  1, in the denominator of L1/L2  (ie L2=0 has p=0)
bool numericalEstimatePoissRatioConf(float lambda1,float lambda2, float alpha, 
			unsigned int nTrials, gsl_rng *r,float &lBound, float &uBound )
{
	vector<float> values;
	values.resize(nTrials);
	using std::isinf;
	using std::isnan;
	if(isinf(lambda1) || isnan(lambda1) ||
		isinf(lambda2) || isnan(lambda2))
		return false;
	
	//will always be zero for l1
	if(lambda1 == 0)
		return false;

	//Compute Z=X/Y
	float vMin,vMax;
	vMin=std::numeric_limits<float>::max();
	vMax=-vMin;
	for(auto &v : values)
	{
		v=gsl_ran_poisson(r,lambda1)/(float)std::max(gsl_ran_poisson(r,lambda2),1u);
		vMin=std::min(v,vMin);
		vMax=std::max(v,vMax);
	}

	if(vMin == std::numeric_limits<float>::max() ||
		vMax == -std::numeric_limits<float>::max() ||
		(fabs(vMin-vMax) < std::numeric_limits<float>::epsilon()))
		return false;

	//Create histogram
	const unsigned int nBins=sqrt(values.size())/2.0f;
	gsl_histogram *h=gsl_histogram_alloc(nBins);
	gsl_histogram_set_ranges_uniform(h,vMin,vMax);
	for(auto &v:values)
		gsl_histogram_increment(h,v);

	//Distribution function. First we compute PDF, then integrate
	vector<float> df;
	df.resize(nBins);
	for(auto ui=0u;ui<nBins;ui++)
		df[ui]=gsl_histogram_get(h,ui);

	gsl_histogram_free(h);

	//Normalise to create PDF
	double dSum=0;;
	for(auto v : df)
		dSum +=v; 
	for(auto &v : df)
		v/=dSum; 

	//OK, so now this gives us the PDF, we have to compute the CDF
	double accum=0;
	for(auto &v  :df)
	{
		accum+=v;
		v=accum;
	}

	//Find upper confidence limit
	//===
	auto alphaIt = std::lower_bound(df.begin(),df.end(),alpha);

	if(alphaIt == df.end())
		return false;

	unsigned int offset = std::distance(df.begin(),alphaIt);
	
	if(offset > 1)
	{
		float dy= df[offset] - df[offset-1];
		float dx= (vMax-vMin)/(float)df.size();

		uBound = (alpha-df[offset-1])*dx/dy + (offset-1)*dx + vMin;
	}
	else
		return false;
	//===

	//Find lower confidence limit
	//===
	alphaIt = std::upper_bound(df.begin(),df.end(),1.0-alpha);
	if(alphaIt==df.begin())
		lBound=0;
	else
	{

		offset = std::distance(df.begin(),alphaIt);
		
		if(offset < df.size()-1)
		{
			float dy= df[offset] - df[offset-1];
			float dx= (vMax-vMin)/(float)df.size();

			lBound = ((1.0-alpha)-df[offset])*dx/dy + offset*dx +vMin;
		}
		else
			return false;
	}
	//===


	ASSERT(lBound<=lambda1/lambda2);
	ASSERT(uBound>=lambda1/lambda2);
	ASSERT(lBound<=uBound);

	return true;
}

//Brute-force Skellam  confidence estimator
bool numericalEstimateSkellamConf(float lambda1,float lambda2, float alpha, 
			unsigned int nTrials, gsl_rng *r,float &lBound, float &uBound )
{
	vector<float> values;
	values.resize(nTrials);

	//Compute Z=X-Y
	float vMin,vMax;
	vMin=std::numeric_limits<float>::max();
	vMax=-vMin;
	for(auto &v : values)
	{
		float v1,v2;
		v1=gsl_ran_poisson(r,lambda1);
		v2=gsl_ran_poisson(r,lambda2);

		v=v1-v2;
		vMin=std::min(v,vMin);
		vMax=std::max(v,vMax);
	}

	//Create histogram
	const unsigned int nBins=sqrt(values.size())/2.0f;
	gsl_histogram *h=gsl_histogram_alloc(nBins);
	gsl_histogram_set_ranges_uniform(h,vMin,vMax);
	for(auto &v:values)
		gsl_histogram_increment(h,v);
	


	//Distribution function. First we compute PDF, then integrate
	vector<float> df;
	df.resize(nBins);
	for(auto ui=0u;ui<nBins;ui++)
		df[ui]=gsl_histogram_get(h,ui);

	gsl_histogram_free(h);

	//Normalise to create PDF
	double dSum=0;;
	for(auto v : df)
		dSum +=v; 
	for(auto &v : df)
		v/=dSum; 

	//OK, so now this gives us the PDF, we have to compute the CDF
	double accum=0;
	for(auto &v  :df)
	{
		accum+=v;
		v=accum;
	}

	//Find upper confidence limit
	//===
	auto alphaIt = std::lower_bound(df.begin(),df.end(),alpha);

	if(alphaIt == df.end())
		return false;

	unsigned int offset = std::distance(df.begin(),alphaIt);
	
	if(offset > 1)
	{
		float dy= df[offset] - df[offset-1];
		float dx= (vMax-vMin)/(float)df.size();

		uBound = (alpha-df[offset-1])*dx/dy + (offset-1)*dx + vMin;
	}
	else
		return false;
	//===

	//Find lower confidence limit
	//===
	alphaIt = std::upper_bound(df.begin(),df.end(),1.0-alpha);
	if(alphaIt==df.begin())
		lBound=0;
	else
	{

		offset = std::distance(df.begin(),alphaIt);
		
		if(offset < df.size()-1)
		{
			float dy= df[offset] - df[offset-1];
			float dx= (vMax-vMin)/(float)df.size();

			lBound = ((1.0-alpha)-df[offset])*dx/dy + offset*dx +vMin;
		}
		else
			return false;
	}
	//===


	ASSERT(lBound<=uBound);

	return true;
}
//Brute-force gaussian ratio confidence estimator
// Biased, as it enforces  >1, in the denominator of L1/L2  (ie L2=0 has p=0)
bool numericalEstimateGaussRatioConf(float mu1,float mu2,float var1, float var2, float alpha, 
			unsigned int nTrials, gsl_rng *r,float &lBound, float &uBound )
{
	vector<float> values;
	values.resize(nTrials);

	//Compute Z=X/Y
	float vMin,vMax;
	vMin=std::numeric_limits<float>::max();
	vMax=-vMin;
	for(auto &v : values)
	{
		v=(gsl_ran_gaussian(r,sqrt(var1))+mu1)/
			std::max(gsl_ran_gaussian(r,sqrt(var2))+mu2,1.0);
		vMin=std::min(v,vMin);
		vMax=std::max(v,vMax);
	}
	
	if(vMin >=vMax)
		return false;
	
	if(vMin == std::numeric_limits<float>::max() ||
		vMax == -std::numeric_limits<float>::max() ||
		(fabs(vMin-vMax) < std::numeric_limits<float>::epsilon()))
		return false;

	//Create histogram
	const unsigned int nBins=sqrt(values.size())/2.0f;
	gsl_histogram *h=gsl_histogram_alloc(nBins);
	gsl_histogram_set_ranges_uniform(h,vMin,vMax);
	for(auto &v:values)
		gsl_histogram_increment(h,v);

	//Distribution function. First we compute PDF, then integrate
	vector<float> df;
	df.resize(nBins);
	for(auto ui=0u;ui<nBins;ui++)
		df[ui]=gsl_histogram_get(h,ui);

	gsl_histogram_free(h);

	//Normalise to create PDF
	double dSum=0;;
	for(auto v : df)
		dSum +=v; 
	for(auto &v : df)
		v/=dSum; 

	//OK, so now this gives us the PDF, we have to compute the CDF
	double accum=0;
	for(auto &v  :df)
	{
		accum+=v;
		v=accum;
	}

	//Find upper confidence limit
	//===
	auto alphaIt = std::lower_bound(df.begin(),df.end(),alpha);

	if(alphaIt == df.end())
		return false;

	unsigned int offset = std::distance(df.begin(),alphaIt);
	
	if(offset > 1)
	{
		float dy= df[offset] - df[offset-1];
		float dx= (vMax-vMin)/(float)df.size();

		uBound = (alpha-df[offset-1])*dx/dy + (offset-1)*dx + vMin;
	}
	else
		return false;
	//===

	//Find lower confidence limit
	//===
	alphaIt = std::upper_bound(df.begin(),df.end(),1.0-alpha);
	if(alphaIt==df.begin())
		lBound=0;
	else
	{

		offset = std::distance(df.begin(),alphaIt);
		
		if(offset < df.size()-1)
		{
			float dy= df[offset] - df[offset-1];
			float dx= (vMax-vMin)/(float)df.size();

			lBound = ((1.0-alpha)-df[offset])*dx/dy + offset*dx +vMin;
		}
		else
			return false;
	}
	//===


	ASSERT(lBound<=mu1/mu2);
	ASSERT(uBound>=mu1/mu2);
	ASSERT(lBound<=uBound);

	return true;
}


struct ZECH_ROOT
{
	float alpha;
	float lambdaBack;
	float observation;
};

//Find the root of zech CDF with a specified alpha, at observed signal+back = "observation"
// and background intensity of "lambdaBack"
double zechRoot(double sigGuess,void *params)
{
	ZECH_ROOT *zr = (ZECH_ROOT*)params;

	//Initial guess for signal
	double denom=1.0/gsl_cdf_poisson_P(zr->observation,zr->lambdaBack);
	double epsilon=0;

	//Add the component for each observation's probability
	for(auto i=0u; i<=zr->observation;i++)
		epsilon+= gsl_ran_poisson_pdf(i,sigGuess+zr->lambdaBack);

	epsilon*=denom;	
	//Convert to confidence
	epsilon=1-epsilon;
	
	//subtract to get confidence we want as root
	return epsilon-zr->alpha;
}

//Obtain the confidence  limits at a specified alpha for zech, using a bisection approach
bool zechConfidenceLimits(float lambdaBack, unsigned int observation, float alpha,
				float &estimate)
{
	ASSERT(lambdaBack > 0);

	//We are working with Equation 3, and the combined distribution W
	// in Zech.


	ZECH_ROOT zr;
	zr.alpha=alpha; 
	zr.observation=observation;
	zr.lambdaBack=lambdaBack;

	gsl_function F;
	F.function = zechRoot;
	F.params=&zr;
	
	float xLow=0;
	float xHi=std::max((float)observation,lambdaBack)*5.0f;
	
	//Use Brent's method
	gsl_root_fsolver *s = gsl_root_fsolver_alloc (gsl_root_fsolver_brent);
	gsl_root_fsolver_set (s, &F, xLow,xHi);


	const unsigned int MAXITS=1000;

	auto it=0,status=0;
	do
	{

		status = gsl_root_fsolver_iterate (s);
		xLow = gsl_root_fsolver_x_lower (s);
		xHi = gsl_root_fsolver_x_upper (s);
		status = gsl_root_test_interval (xLow, xHi,
						 0, 0.01);

		if (status == GSL_SUCCESS)
			break;
	} while(status ==GSL_CONTINUE && it < MAXITS);


	estimate = gsl_root_fsolver_root (s);
	gsl_root_fsolver_free(s);

	//Check for sanity
	if(estimate<0 || isnan(estimate) || it >=MAXITS)
		return false;


	return true;
}


template<class T>
void cumTrapezoid(const vector<T> &x, const vector<T> &vals, vector<T> &res)
{
	ASSERT(vals.size() ==x.size());

	vector<T> sumVals;
	sumVals.resize(vals.size()+1);
	sumVals[0]=0;

	//Perform stable summation
	kahansum(vals.begin(),vals.begin()+vals.size(),sumVals.begin()+1);

	res.resize(x.size());
	res[0]=0;
	for(auto ui=1u;ui<vals.size();ui++)
		res[ui] = 0.5*(sumVals[ui-1] + sumVals[ui])*(x[ui]-x[ui-1]);


}

// https://en.wikipedia.org/wiki/Ratio_distribution#Gaussian_ratio_distribution (4 Apr, 2019) 
// Checked with simple M-C simulation.  The MC simulation seems to converge for trialed values
// as the MC support increases (as we artificially chopped the upper and lower simulation bounds at times
double gauss_ratio_pdf(double x, double muX,double muY, double varX,double varY)
{
	//Original equation is Eqn 1 & 2,
	// From Hinkley, Biometrika, 1969. 56, 3, pp635
	// DOI : 10.1093/biomet/56.3.635
	// We assume rho (cross-correlation coefficient) = 0
	double a = sqrt(x*x/varX + 1.0/varY);
	double b =muX/varX*x + muY/varY;
	double c = muX*muX/varX + muY*muY/varY;
	double lnD = (b*b - c*a*a)/(2.0*a*a);

	double d = exp(lnD);

	double bDivA = b/a;

	double theta1 = 0.5*(1.0+erf(bDivA/sqrt(2.0)));
	double theta2 = 0.5*(1.0+erf(-bDivA/sqrt(2.0)));

	double p1 = b*d/(a*a*a)*1/(sqrt(2*M_PI*varX*varY))*(theta1-theta2);

	double p2 = 1.0/(a*a*M_PI*sqrt(varX*varY))*exp(-c/2.0);

	return p1+p2;
}

/* Problems with numerical stability, owing to very sharp PDF
   	- Either we need to estimate the support bounds better, or we
	  need to integrate the bivariate normal distribution numerically
	  to obtain the CDF semi-analytically
//Note: Future improvement: 
// Hinkley (DOI 10.1093/biomet/56.3.635)  provides what looks like an analytic CDF 
// which has integrals in it. In this version, we analytically obtain the PDF, then numerically
// convert this to a CDF on a fixed support. An analytical expression could enable
// providing an explicit support and bin width unnecessary
bool estimateGaussianRatioConf(float muX, float muY,float varX, float varY,float alpha, 
			float &lBound, float &uBound, unsigned int nBins,float supportBound)
{
	ASSERT(varX >0);
	ASSERT(varY >0);
	ASSERT(alpha >0.5 && alpha <1.0f);

	uBound=std::numeric_limits<float>::max();
	lBound=-uBound;

	//Compute PDF
	vector<double> support;
	support.resize(nBins);
	for(auto ui=0u;ui<support.size();ui++)
		support[ui] = (ui - support.size()/2.0)*supportBound/nBins; 

	vector<double> pdf;
	pdf.resize(support.size());
	for(auto ui=0u;ui<pdf.size();ui++)
		pdf[ui] =  gauss_ratio_pdf(support[ui],muX,muY,varX,varY);

	//Integrate PDF. Numerical drift can cause this to exceed 1
	vector<double> cdf;
	cumTrapezoid(support,pdf,cdf);

	//Normalise CDF
	double cdfMax = *std::max_element(cdf.begin(),cdf.end());
	for(auto &v : cdf )
		v/=cdfMax;

	unsigned int alphaLowIdx=-1;
	unsigned int alphaHighIdx=-1;
	float oneAlpha = 1.0-alpha;

	auto alphaIt = std::lower_bound(cdf.begin(),cdf.end(),alpha);
	if(alphaIt == cdf.end())
		return false;
	alphaHighIdx=std::distance(cdf.begin(),alphaIt);

	alphaIt = std::upper_bound(cdf.begin(),cdf.end(),oneAlpha);
	if(alphaIt == cdf.end())
		return false;
	alphaLowIdx=std::distance(cdf.begin(),alphaIt);

	if(alphaHighIdx == (unsigned int)-1 || alphaLowIdx == (unsigned int) -1 ||
			alphaHighIdx == 0 || alphaLowIdx == (cdf.size()-1)) 
		return false;
	else
	{
		//Linear interpolate
		double dx,dy;	
		dx = support[alphaHighIdx] - support[alphaHighIdx-1];
		dy = cdf[alphaHighIdx] - cdf[alphaHighIdx-1];
		uBound = (alpha-cdf[alphaHighIdx])*dx/dy + support[alphaHighIdx];

		dx = support[alphaLowIdx+1] - support[alphaLowIdx];
		dy = cdf[alphaLowIdx+1] - cdf[alphaLowIdx];
		lBound = (oneAlpha-cdf[alphaLowIdx])*dx/dy + support[alphaLowIdx];

	}

	ASSERT(lBound <=muX/muY);
	ASSERT(uBound >=muX/muY);
	
	ASSERT(lBound<=uBound);

	return true;
}
*/

void poissonConfidenceObservation(float counts, float alpha, float &lBound, float &uBound)
{
	ASSERT(counts >=0);
	ASSERT(alpha>0 && alpha < 1);

	float a = gsl_cdf_chisq_Pinv(alpha,1);


	lBound = counts + a/2.0 - sqrtf(a*(counts+a/4.0));
	uBound = counts + a/2.0 + sqrtf(a*(counts+a/4.0));
/*
	float alphaOne=(1.0f-alpha)/2.0;
	float alphaTwo = 1.0-alphaOne;

	if(!counts)
		lBound=0;
	else
		lBound=gsl_cdf_chisq_Pinv(alphaOne,2.0*observedCounts)/2.0f;

	uBound=gsl_cdf_chisq_Pinv(alphaTwo,2.0*(observedCounts+1))/2.0f;
	*/
}

#ifdef DEBUG

bool runConfidenceTests()
{
	using std::make_pair;
	using std::map;
	using std::pair;
	
	gsl_rng *rng =gsl_rng_alloc(gsl_rng_mt19937);
	auto t = std::chrono::system_clock::now();
	time_t tt  =std::chrono::system_clock::to_time_t(t);
	gsl_rng_set(rng, tt);                  // set seed

	float lBound,uBound;

	//Map (lambda1,lambda2) -> skellam uBound
	map<pair<float,float>, float> skelMap;
	
	//Values at alpha=0.99
	//FIXME: These values were derived using a "hard cut" of the CDF
	// (finding alpha > 0.99, over a limited support).
	// Thus these are an over-estimation
	skelMap[make_pair(100,100)]=141; 
	skelMap[make_pair(1,10)]=15;

	using std::endl;
	using std::cerr;

	for(auto v : skelMap)
	{
		bool returnCode;
		float l1,l2,res;
		l1=v.first.first;
		l2=v.first.second;
		res=v.second;
		returnCode=numericalEstimateSkellamConf(l1+l2,l2,0.99,10000,rng,lBound,uBound);
		
		ASSERT(returnCode);
		ASSERT(fabs(uBound - res)/uBound < 0.5); // This is wide due to earlier overestimation 
	}

	//Test Zech's confidence bounds
	{
		//Map background,observation -> lower,upper confidence pair
		//This is for alpha=(0.01,0.99)
		map<pair<float,float>, pair<float,float> > confMap;
		confMap[make_pair(1,10)]= make_pair(3.659,19.17);

		const float ALPHA=0.99;
		for(auto &v : confMap)
		{
			float low,high;
			TEST(zechConfidenceLimits(v.first.first,
				v.first.second,ALPHA,high),"Zech UB");
			TEST(zechConfidenceLimits(v.first.first,
				v.first.second,1.0-ALPHA,low),"Zech LB");

			auto confPair = v.second;
			TEST(low < confPair.first+0.5f,"Low confidence bound");
			TEST(high > confPair.second-0.5f,"High confidence bound");
		}
	}



	//Poisson confidence values at fixed interval
	//===
	map<unsigned int,pair<float,float> > poissConfVals;
//		Pearson values
	poissConfVals[1] = make_pair(0.1765,5.6649);
	poissConfVals[3] = make_pair(1.0203,8.8212);
	poissConfVals[5] = make_pair(2.1357,11.7058);
/*		"Exact" values
	poissConfVals[1] = make_pair(0.0253,5.5716);
	poissConfVals[3] = make_pair(0.6187,8.7673);
	poissConfVals[5] = make_pair(1.6235,11.6683);
*/

	for( auto v : poissConfVals)
	{
		float lBound,uBound;
		poissonConfidenceObservation(v.first,0.95,lBound,uBound);
		ASSERT(fabs(v.second.first -lBound) < 0.01);
		ASSERT(fabs(v.second.second - uBound) < 0.01);
	}
	//===


	//FIXME: Better unit tests. This simply exercises the 
	// code for crashes, and then runs it repeatedly.
	// No numerical results are checked.
	// also, we do not examine all functions 
	const unsigned int NREPEATS=250;
	for(auto ui=0;ui<NREPEATS;ui++)
	{
		float v1,v2;
		v1=gsl_ran_poisson(rng,360);
		v2=gsl_ran_poisson(rng,125);

		numericalEstimateSkellamConf(v1,v2,0.99,1000,rng,lBound,uBound);

		if(!v2)
			continue;

		numericalEstimatePoissRatioConf(v1,v2,0.95,1000,rng,lBound,uBound);
		numericalEstimateGaussRatioConf(v1,v2,sqrt(v1),sqrt(v2),
						0.95,1000,rng,lBound,uBound);

	}

	gsl_rng_free(rng);
	return true;
}
#endif
