/*!
  \file
  \brief Probability distributions for Exponential Family models.
  \author Vaclav Smidl.

  -----------------------------------
  BDM++ - C++ library for Bayesian Decision Making under Uncertainty

  Using IT++ for numerical operations
  -----------------------------------
*/

#ifndef EF_H
#define EF_H

#include <itpp/itbase.h>
#include "../math/libDC.h"
#include "libBM.h"
#include "../itpp_ext.h"
//#include <std>

using namespace itpp;


//! Global Uniform_RNG
extern Uniform_RNG UniRNG;
//! Global Normal_RNG
extern Normal_RNG NorRNG;
//! Global Gamma_RNG
extern Gamma_RNG GamRNG;

/*!
* \brief General conjugate exponential family posterior density.

* More?...
*/

class eEF : public epdf {
public:
//	eEF() :epdf() {};
	//! default constructor
	eEF ( const RV &rv ) :epdf ( rv ) {};
	//! logarithm of the normalizing constant, \f$\mathcal{I}\f$
	virtual double lognc() const =0;
	//!TODO decide if it is really needed
	virtual void dupdate ( mat &v ) {it_error ( "Not implemented" );};
	//!Evaluate normalized log-probability
	virtual double evalpdflog_nn ( const vec &val ) const{it_error ( "Not implemented" );return 0.0;};
	//!Evaluate normalized log-probability
	virtual double evalpdflog ( const vec &val ) const {return evalpdflog_nn ( val )-lognc();}
	//!Evaluate normalized log-probability for many samples
	virtual vec evalpdflog ( const mat &Val ) const {
		vec x ( Val.cols() );
		for ( int i=0;i<Val.cols();i++ ) {x ( i ) =evalpdflog_nn ( Val.get_col ( i ) ) ;}
		return x-lognc();
	}
	//!Power of the density, used e.g. to flatten the density
	virtual void pow ( double p ) {it_error ( "Not implemented" );};
};

/*!
* \brief Exponential family model.

* More?...
*/

class mEF : public mpdf {

public:
	//! Default constructor
	mEF ( const RV &rv0, const RV &rvc0 ) :mpdf ( rv0,rvc0 ) {};
};

//! Estimator for Exponential family
class BMEF : public BM {
protected:
	//! forgetting factor
	double frg;
	//! cached value of lognc() in the previous step (used in evaluation of \c ll )
	double last_lognc;
public:
	//! Default constructor
	BMEF ( const RV &rv, double frg0=1.0 ) :BM ( rv ), frg ( frg0 ) {}
	//! Copy constructor
	BMEF ( const BMEF &B ) :BM ( B ), frg ( B.frg ), last_lognc ( B.last_lognc ) {}
	//!get statistics from another model
	virtual void set_statistics ( const BMEF* BM0 ) {it_error ( "Not implemented" );};
	//! Weighted update of sufficient statistics (Bayes rule)
	virtual void bayes ( const vec &data, const double w ) {};
	//original Bayes
	void bayes ( const vec &dt );
	//!Flatten the posterior according to the given BMEF (of the same type!)
	virtual void flatten ( const BMEF * B ) {it_error ( "Not implemented" );}
	//!Flatten the posterior as if to keep nu0 data
	virtual void flatten ( double nu0 ) {it_error ( "Not implemented" );}
};

template<class sq_T>
class mlnorm;

/*!
* \brief Gaussian density with positive definite (decomposed) covariance matrix.

* More?...
*/
template<class sq_T>
class enorm : public eEF {
protected:
	//! mean value
	vec mu;
	//! Covariance matrix in decomposed form
	sq_T R;
	//! dimension (redundant from rv.count() for easier coding )
	int dim;
public:
	//!Default constructor
	enorm ( const RV &rv );
	//! Set mean value \c mu and covariance \c R
	void set_parameters ( const vec &mu,const sq_T &R );
	////! tupdate in exponential form (not really handy)
	//void tupdate ( double phi, mat &vbar, double nubar );
	//! dupdate in exponential form (not really handy)
	void dupdate ( mat &v,double nu=1.0 );

	vec sample() const;
	//! TODO is it used?
	mat sample ( int N ) const;
	double eval ( const vec &val ) const ;
	double evalpdflog_nn ( const vec &val ) const;
	double lognc () const;
	vec mean() const {return mu;}
	mlnorm<sq_T>* condition ( const RV &rvn );
	enorm<sq_T>* marginal ( const RV &rv );
//Access methods
	//! returns a pointer to the internal mean value. Use with Care!
	vec& _mu() {return mu;}

	//! access function
	void set_mu ( const vec mu0 ) { mu=mu0;}

	//! returns pointers to the internal variance and its inverse. Use with Care!
	sq_T& _R() {return R;}

	//! access method
//	mat getR () {return R.to_mat();}
};

/*!
* \brief Gauss-inverse-Wishart density stored in LD form

* For \f$p\f$-variate densities, given rv.count() should be \f$p\times\f$ V.rows().
*
*/
class egiw : public eEF {
protected:
	//! Extended information matrix of sufficient statistics
	ldmat V;
	//! Number of data records (degrees of freedom) of sufficient statistics
	double nu;
	//! Dimension of the output
	int xdim;
	//! Dimension of the regressor
	int nPsi;
public:
	//!Default constructor, assuming
	egiw ( RV rv, mat V0, double nu0 ) : eEF ( rv ), V ( V0 ), nu ( nu0 ) {
		xdim = rv.count() /V.rows();
		it_assert_debug ( rv.count() ==xdim*V.rows(),"Incompatible V0." );
		nPsi = V.rows()-xdim;
	}
	//!Full constructor for V in ldmat form
	egiw ( RV rv, ldmat V0, double nu0 ) : eEF ( rv ), V ( V0 ), nu ( nu0 ) {
		xdim = rv.count() /V.rows();
		it_assert_debug ( rv.count() ==xdim*V.rows(),"Incompatible V0." );
		nPsi = V.rows()-xdim;
	}

	vec sample() const;
	vec mean() const;
	void mean_mat ( mat &M, mat&R ) const;
	//! In this instance, val= [theta, r]. For multivariate instances, it is stored columnwise val = [theta_1 theta_2 ... r_1 r_2 ]
	double evalpdflog_nn ( const vec &val ) const;
	double lognc () const;

	//Access
	//! returns a pointer to the internal statistics. Use with Care!
	ldmat& _V() {return V;}
	//! returns a pointer to the internal statistics. Use with Care!
	double& _nu() {return nu;}
	void pow ( double p );
};

/*! \brief Dirichlet posterior density

Continuous Dirichlet density of \f$n\f$-dimensional variable \f$x\f$
\f[
f(x|\beta) = \frac{\Gamma[\gamma]}{\prod_{i=1}^{n}\Gamma(\beta_i)} \prod_{i=1}^{n}x_i^{\beta_i-1}
\f]
where \f$\gamma=\sum_i \beta_i\f$.
*/
class eDirich: public eEF {
protected:
	//!sufficient statistics
	vec beta;
public:
	//!Default constructor
	eDirich ( const RV &rv, const vec &beta0 ) : eEF ( rv ),beta ( beta0 ) {it_assert_debug ( rv.count() ==beta.length(),"Incompatible statistics" ); };
	//! Copy constructor
	eDirich ( const eDirich &D0 ) : eEF ( D0.rv ),beta ( D0.beta ) {};
	vec sample() const {it_error ( "Not implemented" );return vec_1 ( 0.0 );};
	vec mean() const {return beta/sum ( beta );};
	//! In this instance, val is ...
	double evalpdflog_nn ( const vec &val ) const {return ( beta-1 ) *log ( val );};
	double lognc () const {
		double gam=sum ( beta );
		double lgb=0.0;
		for ( int i=0;i<beta.length();i++ ) {lgb+=lgamma ( beta ( i ) );}
		return lgb-lgamma ( gam );
	};
	//!access function
	vec& _beta()  {return beta;}
	//!Set internal parameters
	void set_parameters ( const vec &beta0 ) {
		if ( beta0.length() !=beta.length() ) {
			it_assert_debug ( rv.length() ==1,"Undefined" );
			rv.set_size ( 0,beta0.length() );
		}
		beta= beta0;
	}
};

//! Estimator for Multinomial density
class multiBM : public BMEF {
protected:
	//! Conjugate prior and posterior
	eDirich est;
	vec &beta;
public:
	//!Default constructor
	multiBM ( const RV &rv, const vec beta0 ) : BMEF ( rv ),est ( rv,beta0 ),beta ( est._beta() ) {last_lognc=est.lognc();}
	//!Copy constructor
	multiBM ( const multiBM &B ) : BMEF ( B ),est ( rv,B.beta ),beta ( est._beta() ) {}

	void set_statistics ( const BM* mB0 ) {const multiBM* mB=dynamic_cast<const multiBM*> ( mB0 ); beta=mB->beta;}
	void bayes ( const vec &dt ) {
		if ( frg<1.0 ) {beta*=frg;last_lognc=est.lognc();}
		beta+=dt;
		if ( evalll ) {ll=est.lognc()-last_lognc;}
	}
	double logpred ( const vec &dt ) const {
		eDirich pred ( est );
		vec &beta = pred._beta();

		double lll;
		if ( frg<1.0 )
			{beta*=frg;lll=pred.lognc();}
		else
			if ( evalll ) {lll=last_lognc;}
			else{lll=pred.lognc();}

		beta+=dt;
		return pred.lognc()-lll;
	}
	void flatten ( const BMEF* B ) {
		const eDirich* E=dynamic_cast<const eDirich*> ( B );
		// sum(beta) should be equal to sum(B.beta)
		const vec &Eb=const_cast<eDirich*> ( E )->_beta();
		est.pow ( sum ( beta ) /sum ( Eb ) );
		if ( evalll ) {last_lognc=est.lognc();}
	}
	const epdf& _epdf() const {return est;};
	void set_parameters ( const vec &beta0 ) {
		est.set_parameters ( beta0 );
		rv = est._rv();
		if ( evalll ) {last_lognc=est.lognc();}
	}
};

/*!
 \brief Gamma posterior density

 Multivariate Gamma density as product of independent univariate densities.
 \f[
 f(x|\alpha,\beta) = \prod f(x_i|\alpha_i,\beta_i)
 \f]
*/

class egamma : public eEF {
protected:
	//! Vector \f$\alpha\f$
	vec alpha;
	//! Vector \f$\beta\f$
	vec beta;
public :
	//! Default constructor
	egamma ( const RV &rv ) :eEF ( rv ) {};
	//! Sets parameters
	void set_parameters ( const vec &a, const vec &b ) {alpha=a,beta=b;};
	vec sample() const;
	//! TODO: is it used anywhere?
//	mat sample ( int N ) const;
	double evalpdflog ( const vec &val ) const;
	double lognc () const;
	//! Returns poiter to alpha and beta. Potentially dengerous: use with care!
	void _param ( vec* &a, vec* &b ) {a=&alpha;b=&beta;};
	vec mean() const {vec pom ( alpha ); pom/=beta; return pom;}
};
/*
//! Weighted mixture of epdfs with external owned components.
class emix : public epdf {
protected:
	int n;
	vec &w;
	Array<epdf*> Coms;
public:
//! Default constructor
	emix ( const RV &rv, vec &w0): epdf(rv), n(w0.length()), w(w0), Coms(n) {};
	void set_parameters( int &i, double wi, epdf* ep){w(i)=wi;Coms(i)=ep;}
	vec mean(){vec pom; for(int i=0;i<n;i++){pom+=Coms(i)->mean()*w(i);} return pom;};
	vec sample() {it_error ( "Not implemented" );return 0;}
};
*/

//!  Uniform distributed density on a rectangular support

class euni: public epdf {
protected:
//! lower bound on support
	vec low;
//! upper bound on support
	vec high;
//! internal
	vec distance;
//! normalizing coefficients
	double nk;
//! cache of log( \c nk )
	double lnk;
public:
	//! Defualt constructor
	euni ( const RV rv ) :epdf ( rv ) {}
	double eval ( const vec &val ) const  {return nk;}
	double evalpdflog ( const vec &val ) const  {return lnk;}
	vec sample() const {
		vec smp ( rv.count() );
#pragma omp critical
		UniRNG.sample_vector ( rv.count(),smp );
		return low+elem_mult ( distance,smp );
	}
	//! set values of \c low and \c high
	void set_parameters ( const vec &low0, const vec &high0 ) {
		distance = high0-low0;
		it_assert_debug ( min ( distance ) >0.0,"bad support" );
		low = low0;
		high = high0;
		nk = prod ( 1.0/distance );
		lnk = log ( nk );
	}
	vec mean() const {vec pom=high; pom-=low; pom/=2.0; return pom;}
};


/*!
 \brief Normal distributed linear function with linear function of mean value;

 Mean value \f$mu=A*rvc+mu_0\f$.
*/
template<class sq_T>
class mlnorm : public mEF {
	//! Internal epdf that arise by conditioning on \c rvc
	enorm<sq_T> epdf;
	mat A;
	vec mu_const;
	vec& _mu; //cached epdf.mu;
public:
	//! Constructor
	mlnorm (const RV &rv, const RV &rvc );
	//! Set \c A and \c R
	void set_parameters ( const  mat &A, const vec &mu0, const sq_T &R );
	//!Generate one sample of the posterior
	vec samplecond (const vec &cond, double &lik );
	//!Generate matrix of samples of the posterior
	mat samplecond (const vec &cond, vec &lik, int n );
	//! Set value of \c rvc . Result of this operation is stored in \c epdf use function \c _ep to access it.
	void condition (const vec &cond );
};

/*!
\brief  Gamma random walk

Mean value, \f$\mu\f$, of this density is given by \c rvc .
Standard deviation of the random walk is proportional to one \f$k\f$-th the mean.
This is achieved by setting \f$\alpha=k\f$ and \f$\beta=k/\mu\f$.

The standard deviation of the walk is then: \f$\mu/\sqrt(k)\f$.
*/
class mgamma : public mEF {
protected:
	//! Internal epdf that arise by conditioning on \c rvc
	egamma epdf;
	//! Constant \f$k\f$
	double k;
	//! cache of epdf.beta
	vec* _beta;

public:
	//! Constructor
	mgamma ( const RV &rv,const RV &rvc );
	//! Set value of \c k
	void set_parameters ( double k );
	void condition ( const vec &val ) {*_beta=k/val;};
};

/*!
\brief  Gamma random walk around a fixed point

Mean value, \f$\mu\f$, of this density is given by a geometric combination of \c rvc and given fixed point, \f$p\f$. \f$l\f$ is the coefficient of the geometric combimation
\f[ \mu = \mu_{t-1} ^{l} p^{1-l}\f]

Standard deviation of the random walk is proportional to one \f$k\f$-th the mean.
This is achieved by setting \f$\alpha=k\f$ and \f$\beta=k/\mu\f$.

The standard deviation of the walk is then: \f$\mu/\sqrt(k)\f$.
*/
class mgamma_fix : public mgamma {
protected:
	//! parameter l
	double l;
	//! reference vector
	vec refl;
public:
	//! Constructor
	mgamma_fix ( const RV &rv,const RV &rvc ) : mgamma ( rv,rvc ),refl ( rv.count() ) {};
	//! Set value of \c k
	void set_parameters ( double k0 , vec ref0, double l0 ) {
		mgamma::set_parameters ( k0 );
		refl=pow ( ref0,1.0-l0 );l=l0;
	};

	void condition ( const vec &val ) {vec mean=elem_mult ( refl,pow ( val,l ) ); *_beta=k/mean;};
};

//! Switch between various resampling methods.
enum RESAMPLING_METHOD { MULTINOMIAL = 0, STRATIFIED = 1, SYSTEMATIC = 3 };
/*!
\brief Weighted empirical density

Used e.g. in particle filters.
*/
class eEmp: public epdf {
protected :
	//! Number of particles
	int n;
	//! Sample weights \f$w\f$
	vec w;
	//! Samples \f$x^{(i)}, i=1..n\f$
	Array<vec> samples;
public:
	//! Default constructor
	eEmp ( const RV &rv0 ,int n0 ) :epdf ( rv0 ),n ( n0 ),w ( n ),samples ( n ) {};
	//! Set samples and weights
	void set_parameters ( const vec &w0, const epdf* pdf0 );
	//! Set sample
	void set_samples ( const epdf* pdf0 );
	//! Potentially dangerous, use with care.
	vec& _w()  {return w;};
	//! access function
	Array<vec>& _samples() {return samples;};
	//! Function performs resampling, i.e. removal of low-weight samples and duplication of high-weight samples such that the new samples represent the same density.
	ivec resample ( RESAMPLING_METHOD method = SYSTEMATIC );
	//! inherited operation : NOT implemneted
	vec sample() const {it_error ( "Not implemented" );return 0;}
	//! inherited operation : NOT implemneted
	double evalpdflog ( const vec &val ) const {it_error ( "Not implemented" );return 0.0;}
	vec mean() const {
		vec pom=zeros ( rv.count() );
		for ( int i=0;i<n;i++ ) {pom+=samples ( i ) *w ( i );}
		return pom;
	}
};


////////////////////////

template<class sq_T>
enorm<sq_T>::enorm ( const RV &rv ) :eEF ( rv ), mu ( rv.count() ),R ( rv.count() ),dim ( rv.count() ) {};

template<class sq_T>
void enorm<sq_T>::set_parameters ( const vec &mu0, const sq_T &R0 ) {
//Fixme test dimensions of mu0 and R0;
	mu = mu0;
	R = R0;
};

template<class sq_T>
void enorm<sq_T>::dupdate ( mat &v, double nu ) {
	//
};

// template<class sq_T>
// void enorm<sq_T>::tupdate ( double phi, mat &vbar, double nubar ) {
// 	//
// };

template<class sq_T>
vec enorm<sq_T>::sample() const {
	vec x ( dim );
	NorRNG.sample_vector ( dim,x );
	vec smp = R.sqrt_mult ( x );

	smp += mu;
	return smp;
};

template<class sq_T>
mat enorm<sq_T>::sample ( int N ) const {
	mat X ( dim,N );
	vec x ( dim );
	vec pom;
	int i;

	for ( i=0;i<N;i++ ) {
		NorRNG.sample_vector ( dim,x );
		pom = R.sqrt_mult ( x );
		pom +=mu;
		X.set_col ( i, pom );
	}

	return X;
};

template<class sq_T>
double enorm<sq_T>::eval ( const vec &val ) const {
	double pdfl,e;
	pdfl = evalpdflog ( val );
	e = exp ( pdfl );
	return e;
};

template<class sq_T>
double enorm<sq_T>::evalpdflog_nn ( const vec &val ) const {
	// 1.83787706640935 = log(2pi)
	return  -0.5* ( R.invqform ( mu-val ) );// - lognc();
};

template<class sq_T>
inline double enorm<sq_T>::lognc () const {
	// 1.83787706640935 = log(2pi)
	return 0.5* ( R.cols() * 1.83787706640935 +R.logdet() );
};

template<class sq_T>
mlnorm<sq_T>::mlnorm ( const RV &rv0, const RV &rvc0 ) :mEF ( rv0,rvc0 ),epdf ( rv0 ),A ( rv0.count(),rv0.count() ),_mu ( epdf._mu() ) {
	ep =&epdf;
}

template<class sq_T>
void mlnorm<sq_T>::set_parameters ( const mat &A0, const vec &mu0, const sq_T &R0 ) {
	epdf.set_parameters ( zeros ( rv.count() ),R0 );
	A = A0;
	mu_const = mu0;
}

template<class sq_T>
vec mlnorm<sq_T>::samplecond (const  vec &cond, double &lik ) {
	this->condition ( cond );
	vec smp = epdf.sample();
	lik = epdf.eval ( smp );
	return smp;
}

template<class sq_T>
mat mlnorm<sq_T>::samplecond (const vec &cond, vec &lik, int n ) {
	int i;
	int dim = rv.count();
	mat Smp ( dim,n );
	vec smp ( dim );
	this->condition ( cond );

	for ( i=0; i<n; i++ ) {
		smp = epdf.sample();
		lik ( i ) = epdf.eval ( smp );
		Smp.set_col ( i ,smp );
	}

	return Smp;
}

template<class sq_T>
void mlnorm<sq_T>::condition (const vec &cond ) {
	_mu = A*cond + mu_const;
//R is already assigned;
}

template<class sq_T>
enorm<sq_T>* enorm<sq_T>::marginal ( const RV &rvn ) {
	ivec irvn = rvn.dataind ( rv );

	sq_T Rn ( R,irvn );
	enorm<sq_T>* tmp = new enorm<sq_T>( rvn );
	tmp->set_parameters ( mu ( irvn ), Rn );
	return tmp;
}

template<class sq_T>
mlnorm<sq_T>* enorm<sq_T>::condition ( const RV &rvn ) {

	RV rvc = rv.subt ( rvn );
	it_assert_debug ( ( rvc.count() +rvn.count() ==rv.count() ),"wrong rvn" );
	//Permutation vector of the new R
	ivec irvn = rvn.dataind ( rv );
	ivec irvc = rvc.dataind ( rv );
	ivec perm=concat ( irvn , irvc );
	sq_T Rn ( R,perm );

	//fixme - could this be done in general for all sq_T?
	mat S=R.to_mat();
	//fixme
	int n=rvn.count()-1;
	int end=R.rows()-1;
	mat S11 = S.get ( 0,n, 0, n );
	mat S12 = S.get ( rvn.count(), end, 0, n );
	mat S22 = S.get ( rvn.count(), end, rvn.count(), end );

	vec mu1 = mu ( irvn );
	vec mu2 = mu ( irvc );
	mat A=S12*inv ( S22 );
	sq_T R_n ( S11 - A *S12.T() );

	mlnorm<sq_T>* tmp=new mlnorm<sq_T> ( rvn,rvc );

	tmp->set_parameters ( A,mu1-A*mu2,R_n );
	return tmp;
}

///////////


#endif //EF_H