#include <math.h>

#include <itpp/base/bessel.h>
#include "exp_family.h"

namespace bdm {

Uniform_RNG UniRNG;
Normal_RNG NorRNG;
Gamma_RNG GamRNG;

using std::cout;

	///////////

void BMEF::bayes( const vec &yt, const vec &cond ) {
	this->bayes_weighted ( yt, cond, 1.0 );
};

void egiw::set_parameters (int dimx0, ldmat V0, double nu0) {
	dimx = dimx0;
	nPsi = V0.rows() - dimx;
	dim = dimx * (dimx + nPsi); // size(R) + size(Theta)
	
	V = V0;
	if (nu0 < 0) {
		nu = 0.1 + nPsi + 2 * dimx + 2; // +2 assures finite expected value of R
		// terms before that are sufficient for finite normalization
	} else {
		nu = nu0;
	}
}

vec egiw::sample() const {
	mat M;
	chmat R;
	sample_mat(M,R);
	
	return concat (cvectorize(M),cvectorize(R.to_mat()));
}

mat egiw::sample_mat(int n) const {
	// TODO - correct approach - convert to product of norm * Wishart
	mat M;
	ldmat Vz;
	ldmat Lam;
	factorize(M,Vz,Lam);
	
	chmat ChLam(Lam.to_mat());
	chmat iChLam;
	ChLam.inv(iChLam);
	
	eWishartCh Omega; //inverse Wishart, result is R, 
	Omega.set_parameters(iChLam,nu-2*nPsi-dimx); // 2*nPsi is there to match numercial simulations - check if analytically correct
	
	mat OmChi;
	mat Z(M.rows(),M.cols());

	mat Mi;
	mat RChiT;
	mat tmp(dimension(), n);
	for (int i=0; i<n;i++){
		OmChi=Omega.sample_mat();
		RChiT=inv(OmChi);
		Z=randn(M.rows(), M.cols());
		Mi = M + RChiT * Z * inv(Vz._L().T() *diag(sqrt(Vz._D())));
	
		tmp.set_col(i,concat (cvectorize(Mi),cvectorize(RChiT*RChiT.T())));
	}
	return tmp;
}

void egiw::sample_mat(mat &Mi, chmat &Ri)const{
	
	// TODO - correct approach - convert to product of norm * Wishart
	mat M;
	ldmat Vz;
	ldmat Lam;
	factorize(M,Vz,Lam);
	
	chmat Ch;
	Ch.setCh(Lam._L()*diag(sqrt(Lam._D())));
	chmat iCh;
	Ch.inv(iCh);
	
	eWishartCh Omega; //inverse Wishart, result is R, 
	Omega.set_parameters(iCh,nu-2*nPsi-dimx); // 2*nPsi is there to match numercial simulations - check if analytically correct
	
	chmat Omi;
	Omi.setCh(Omega.sample_mat());
	
	mat Z=randn(M.rows(), M.cols());
	Mi = M+ Omi._Ch() * Z * inv(Vz._L()*diag(sqrt(Vz._D())));
	Omi.inv(Ri);
}

double egiw::evallog_nn ( const vec &val ) const {
	int vend = val.length() - 1;

	if ( dimx == 1 ) { //same as the following, just quicker.
		double r = val ( vend ); //last entry!
		if ( r < 0 ) return -inf;
		vec Psi ( nPsi + dimx );
		Psi ( 0 ) = -1.0;
		Psi.set_subvector ( 1, val ( 0, vend - 1 ) ); // fill the rest

		double Vq = V.qform ( Psi );
		return -0.5* ( nu*log ( r ) + Vq / r );
	} else {
		mat Th = reshape ( val ( 0, nPsi * dimx - 1 ), nPsi, dimx );
		fsqmat R ( reshape ( val ( nPsi*dimx, vend ), dimx, dimx ) );
		double ldetR = R.logdet();
		if ( ldetR ) return -inf;
		mat Tmp = concat_vertical ( -eye ( dimx ), Th );
		fsqmat iR ( dimx );
		R.inv ( iR );

		return -0.5* ( nu*ldetR + trace ( iR.to_mat() *Tmp.T() *V.to_mat() *Tmp ) );
	}
}

double egiw::lognc() const {
	const vec& D = V._D();

	double m = nu - nPsi - dimx - 1;
#define	log2  0.693147180559945286226763983
#define	logpi 1.144729885849400163877476189
#define log2pi 1.83787706640935
#define Inf std::numeric_limits<double>::infinity()

	double nkG = 0.5 * dimx * ( -nPsi * log2pi + sum ( log ( D ( dimx, D.length() - 1 ) ) ) );
	// temporary for lgamma in Wishart
	double lg = 0;
	for ( int i = 0; i < dimx; i++ ) {
		lg += lgamma ( 0.5 * ( m - i ) );
	}

	double nkW = 0.5 * ( m * sum ( log ( D ( 0, dimx - 1 ) ) ) ) \
	             - 0.5 * dimx * ( m * log2 + 0.5 * ( dimx - 1 ) * log2pi )  - lg;

//	bdm_assert_debug ( ( ( -nkG - nkW ) > -Inf ) && ( ( -nkG - nkW ) < Inf ), "ARX improper" );
	if (-nkG - nkW==Inf){
		cout << "??" <<endl;
	}
	return -nkG - nkW;
}

vec egiw::est_theta() const {
	if ( dimx == 1 ) {
		const mat &L = V._L();
		int end = L.rows() - 1;

		mat iLsub = ltuinv ( L ( dimx, end, dimx, end ) );

		vec L0 = L.get_col ( 0 );

		return iLsub * L0 ( 1, end );
	} else {
		bdm_error ( "ERROR: est_theta() not implemented for dimx>1" );
		return vec();
	}
}

void egiw::factorize(mat &M, ldmat &Vz, ldmat &Lam) const{	
	const mat &L = V._L();
	const vec &D = V._D();
	int end = L.rows() - 1;
	
	Vz=ldmat(L ( dimx, end, dimx, end ), D(dimx,end));
	mat iLsub = ltuinv ( Vz._L());
	// set mean value
	mat Lpsi = L ( dimx, end, 0, dimx - 1 );
	M = iLsub * Lpsi;
	
	Lam =ldmat ( L (0, dimx-1, 0, dimx-1 ), D (0, dimx-1 ) );  //exp val of R
	 if (1){ // test with Peterka
		 mat VF=V.to_mat();
		 mat Vf=VF(0,dimx-1,0, dimx-1);
		 mat Vzf = VF(dimx,end,0,dimx-1);
		 mat VZ = VF(dimx,end,dimx,end);
		 
		 mat Lam2 = Vf-Vzf.T()*inv(VZ)*Vzf;
	 }
}

ldmat egiw::est_theta_cov() const {
	if ( dimx == 1 ) {
		const mat &L = V._L();
		const vec &D = V._D();
		int end = D.length() - 1;

		mat Lsub = L ( 1, end, 1, end );
//		mat Dsub = diag ( D ( 1, end ) );

		ldmat LD(inv(Lsub).T(), 1.0/D(1,end));
		return LD;

	} else {
		bdm_error ( "ERROR: est_theta_cov() not implemented for dimx>1" );
		return ldmat();
	}

}

vec egiw::mean() const {

	if ( dimx == 1 ) {
		const vec &D = V._D();
		int end = D.length() - 1;

		vec m ( dim );
		m.set_subvector ( 0, est_theta() );
		m ( end ) = D ( 0 ) / ( nu - nPsi - 2 * dimx - 2 );
		return m;
	} else {
		mat M;
		mat R;
		mean_mat ( M, R );
		return concat( cvectorize (  M),cvectorize( R ) );
	}

}

vec egiw::variance() const {
	int l = V.rows();
	// cut out rest of lower-right part of V
	const ldmat tmp ( V, linspace ( dimx, l - 1 ) );
	// invert it
	ldmat itmp ( l );
	tmp.inv ( itmp );
	
	// following Wikipedia notation
	// m=nu-nPsi-dimx-1, p=dimx 
	double mp1p=nu-nPsi-2*dimx; // m-p+1
	double mp1m=mp1p-2;         // m-p-1
	
	if ( dimx == 1 ) {
		double cove = V._D() ( 0 ) / mp1m ;

		vec var ( l );
		var.set_subvector ( 0, diag ( itmp.to_mat() ) *cove );
		var ( l - 1 ) = cove * cove / (mp1m-2);
		return var;
	} else {
		ldmat Vll( V, linspace(0,dimx-1)); // top-left part of V
		mat Y=Vll.to_mat();
		mat varY(Y.rows(), Y.cols()); 
		
		double denom = (mp1p-1)*mp1m*mp1m*(mp1m-2);         // (m-p)(m-p-1)^2(m-p-3)
		
		int i,j;
		for ( i=0; i<Y.rows(); i++){
			for ( j=0; j<Y.cols(); j++){
				varY(i,j) = (mp1p*Y(i,j)*Y(i,j) + mp1m * Y(i,i)* Y(j,j)) /denom;
			}
		}
		vec mean_dR = diag(Y)/mp1m; // corresponds to cove
		vec var_th=diag ( itmp.to_mat() );
		vec var_Th ( mean_dR.length()*var_th.length() );
		// diagonal of diag(mean_dR) \kron diag(var_th)
		for (int i=0; i<mean_dR.length(); i++){
			var_Th.set_subvector(i*var_th.length(), var_th*mean_dR(i));  
		}
		
		return concat(var_Th, cvectorize(varY));
	}
}

void egiw::mean_mat ( mat &M, mat&R ) const {
	const mat &L = V._L();
	const vec &D = V._D();
	int end = L.rows() - 1;

	ldmat ldR ( L ( 0, dimx - 1, 0, dimx - 1 ), D ( 0, dimx - 1 ) / ( nu - nPsi - 2*dimx - 2 ) ); //exp val of R
	mat iLsub = ltuinv ( L ( dimx, end, dimx, end ) );

	// set mean value
	mat Lpsi = L ( dimx, end, 0, dimx - 1 );
	M = iLsub * Lpsi;
	R = ldR.to_mat()  ;
}

vec egamma::sample() const {
	vec smp ( dim );
	int i;

	for ( i = 0; i < dim; i++ ) {
		if ( beta ( i ) > std::numeric_limits<double>::epsilon() ) {
			GamRNG.setup ( alpha ( i ), beta ( i ) );
		} else {
			GamRNG.setup ( alpha ( i ), std::numeric_limits<double>::epsilon() );
		}
#pragma omp critical
		smp ( i ) = GamRNG();
	}

	return smp;
}

// mat egamma::sample ( int N ) const {
// 	mat Smp ( rv.count(),N );
// 	int i,j;
//
// 	for ( i=0; i<rv.count(); i++ ) {
// 		GamRNG.setup ( alpha ( i ),beta ( i ) );
//
// 		for ( j=0; j<N; j++ ) {
// 			Smp ( i,j ) = GamRNG();
// 		}
// 	}
//
// 	return Smp;
// }

double egamma::evallog ( const vec &val ) const {
	double res = 0.0; //the rest will be added
	int i;

	if ( any ( val <= 0. ) ) return -inf;
	if ( any ( beta <= 0. ) ) return -inf;
	for ( i = 0; i < dim; i++ ) {
		res += ( alpha ( i ) - 1 ) * std::log ( val ( i ) ) - beta ( i ) * val ( i );
	}
	double tmp = res - lognc();;
	bdm_assert_debug ( std::isfinite ( tmp ), "Infinite value" );
	return tmp;
}

double egamma::lognc() const {
	double res = 0.0; //will be added
	int i;

	for ( i = 0; i < dim; i++ ) {
		res += lgamma ( alpha ( i ) ) - alpha ( i ) * std::log ( beta ( i ) ) ;
	}

	return res;
}

void mgamma::set_parameters ( double k0, const vec &beta0 ) {
	k = k0;
	iepdf.set_parameters ( k * ones ( beta0.length() ), beta0 );
	dimc = iepdf.dimension();
	dim = iepdf.dimension();
}

void eEmp::resample ( ivec &ind, RESAMPLING_METHOD method ) {
	ind = zeros_i ( n );
	ivec N_babies = zeros_i ( n );
	vec cumDist = cumsum ( w );
	vec u ( n );
	int i, j, parent;
	double u0;

	switch ( method ) {
	case MULTINOMIAL:
		u ( n - 1 ) = pow ( UniRNG.sample(), 1.0 / n );

		for ( i = n - 2; i >= 0; i-- ) {
			u ( i ) = u ( i + 1 ) * pow ( UniRNG.sample(), 1.0 / ( i + 1 ) );
		}

		break;

	case STRATIFIED:

		for ( i = 0; i < n; i++ ) {
			u ( i ) = ( i + UniRNG.sample() ) / n;
		}

		break;

	case SYSTEMATIC:
		u0 = UniRNG.sample();

		for ( i = 0; i < n; i++ ) {
			u ( i ) = ( i + u0 ) / n;
		}

		break;

	default:
		bdm_error ( "PF::resample(): Unknown resampling method" );
	}

	// U is now full
	j = 0;

	for ( i = 0; i < n; i++ ) {
		while ( u ( i ) > cumDist ( j ) ) j++;

		N_babies ( j ) ++;
	}
	// We have assigned new babies for each Particle
	// Now, we fill the resulting index such that:
	// * particles with at least one baby should not move *
	// This assures that reassignment can be done inplace;

	// find the first parent;
	parent = 0;
	while ( N_babies ( parent ) == 0 ) parent++;

	// Build index
	for ( i = 0; i < n; i++ ) {
		if ( N_babies ( i ) > 0 ) {
			ind ( i ) = i;
			N_babies ( i ) --; //this index was now replicated;
		} else {
			// test if the parent has been fully replicated
			// if yes, find the next one
			while ( ( N_babies ( parent ) == 0 ) || ( N_babies ( parent ) == 1 && parent > i ) ) parent++;

			// Replicate parent
			ind ( i ) = parent;

			N_babies ( parent ) --; //this index was now replicated;
		}

	}

// copy the internals according to ind
	for ( i = 0; i < n; i++ ) {
		if ( ind ( i ) != i ) {
			samples ( i ) = samples ( ind ( i ) );
		}
		w ( i ) = 1.0 / n;
	}
}

void eEmp::set_statistics ( const vec &w0, const epdf &epdf0 ) {
	dim = epdf0.dimension();
	w = w0;
	w /= sum ( w0 );//renormalize
	n = w.length();
	samples.set_size ( n );

	for ( int i = 0; i < n; i++ ) {
		samples ( i ) = epdf0.sample();
	}
}

void eEmp::set_samples ( const epdf* epdf0 ) {
	w = 1;
	w /= sum ( w );//renormalize

	for ( int i = 0; i < n; i++ ) {
		samples ( i ) = epdf0->sample();
	}
}

void migamma_ref::from_setting ( const Setting &set ) {
	vec ref;
	UI::get ( ref, set, "ref" , UI::compulsory );
	set_parameters ( set["k"], ref, set["l"] );
}

void mlognorm::from_setting ( const Setting &set ) {
	vec mu0;
	UI::get ( mu0, set, "mu0", UI::compulsory );
	set_parameters ( mu0.length(), set["k"] );
	condition ( mu0 );
}

};