/*!
  \file
  \brief Bayesian Filtering for linear Gaussian models (Kalman Filter) and extensions
  \author Vaclav Smidl.

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

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

#ifndef KF_H
#define KF_H


#include "../stat/libFN.h"
#include "../stat/libEF.h"
#include "../math/chmat.h"

namespace bdm {

/*!
* \brief Basic Kalman filter with full matrices (education purpose only)! Will be deleted soon!
*/

class KalmanFull {
protected:
	int dimx, dimy, dimu;
	mat A, B, C, D, R, Q;

	//cache
	mat _Pp, _Ry, _iRy, _K;
public:
	//posterior
	//! Mean value of the posterior density
	vec mu;
	//! Variance of the posterior density
	mat P;

	bool evalll;
	double ll;
public:
	//! Full constructor
	KalmanFull ( mat A, mat B, mat C, mat D, mat R, mat Q, mat P0, vec mu0 );
	//! Here dt = [yt;ut] of appropriate dimensions
	void bayes ( const vec &dt );
	//! print elements of KF
	friend std::ostream &operator<< ( std::ostream &os, const KalmanFull &kf );
	//! For EKFfull;
	KalmanFull() {};
};


/*!
* \brief Kalman filter with covariance matrices in square root form.

Parameter evolution model:\f[ x_t = A x_{t-1} + B u_t + Q^{1/2} e_t \f]
Observation model: \f[ y_t = C x_{t-1} + C u_t + Q^{1/2} w_t. \f]
Where $e_t$ and $w_t$ are independent vectors Normal(0,1)-distributed disturbances.
*/
template<class sq_T>

class Kalman : public BM {
protected:
	//! Indetifier of output rv
	RV rvy;
	//! Indetifier of exogeneous rv
	RV rvu;
	//! cache of rv.count()
	int dimx;
	//! cache of rvy.count()
	int dimy;
	//! cache of rvu.count()
	int dimu;
	//! Matrix A
	mat A;
	//! Matrix B
	mat B;
	//! Matrix C
	mat C;
	//! Matrix D
	mat D;
	//! Matrix Q in square-root form
	sq_T Q;
	//! Matrix R in square-root form
	sq_T R;

	//!posterior density on $x_t$
	enorm<sq_T> est;
	//!preditive density on $y_t$
	enorm<sq_T> fy;

	//! placeholder for Kalman gain
	mat _K;
	//! cache of fy.mu
	vec& _yp;
	//! cache of fy.R
	sq_T& _Ry;
	//!cache of est.mu
	vec& _mu;
	//!cache of est.R
	sq_T& _P;

public:
	//! Default constructor
	Kalman ( );
	//! Copy constructor
	Kalman ( const Kalman<sq_T> &K0 );
	//! Set parameters with check of relevance
	void set_parameters ( const mat &A0,const mat &B0,const mat &C0,const mat &D0,const sq_T &Q0,const sq_T &R0 );
	//! Set estimate values, used e.g. in initialization.
	void set_est ( const vec &mu0, const sq_T &P0 ) {
		sq_T pom ( dimy );
		est.set_parameters ( mu0,P0 );
		P0.mult_sym ( C,pom );
		fy.set_parameters ( C*mu0, pom );
	};

	//! Here dt = [yt;ut] of appropriate dimensions
	void bayes ( const vec &dt );
	//!access function
	const epdf& posterior() const {return est;}
	const enorm<sq_T>* _e() const {return &est;}
	//!access function
	mat& __K() {return _K;}
	//!access function
	vec _dP() {return _P->getD();}
};

/*! \brief Kalman filter in square root form

Trivial example:
\include kalman_simple.cpp

*/
class KalmanCh : public Kalman<chmat> {
protected:
//! pre array (triangular matrix)
	mat preA;
//! post array (triangular matrix)
	mat postA;

public:
	//! copy constructor
	BM* _copy_() const {
		KalmanCh* K=new KalmanCh;
		K->set_parameters ( A,B,C,D,Q,R );
		K->set_statistics ( _mu,_P );
		return K;
	}
	//! Set parameters with check of relevance
	void set_parameters ( const mat &A0,const mat &B0,const mat &C0,const mat &D0,const chmat &Q0,const chmat &R0 );
	void set_statistics ( const vec &mu0, const chmat &P0 ) {
		est.set_parameters ( mu0,P0 );
	};


	/*!\brief  Here dt = [yt;ut] of appropriate dimensions

	The following equality hold::\f[
	\left[\begin{array}{cc}
	R^{0.5}\\
	P_{t|t-1}^{0.5}C' & P_{t|t-1}^{0.5}CA'\\
	& Q^{0.5}\end{array}\right]<\mathrm{orth.oper.}>=\left[\begin{array}{cc}
	R_{y}^{0.5} & KA'\\
	& P_{t+1|t}^{0.5}\\
	\\\end{array}\right]\f]

	Thus this object evaluates only predictors! Not filtering densities.
	*/
	void bayes ( const vec &dt );
};

/*!
\brief Extended Kalman Filter in full matrices

An approximation of the exact Bayesian filter with Gaussian noices and non-linear evolutions of their mean.
*/
class EKFfull : public KalmanFull, public BM {

	//! Internal Model f(x,u)
	diffbifn* pfxu;
	//! Observation Model h(x,u)
	diffbifn* phxu;

	enorm<fsqmat> E;
public:
	//! Default constructor
	EKFfull ( );
	//! Set nonlinear functions for mean values and covariance matrices.
	void set_parameters ( diffbifn* pfxu, diffbifn* phxu, const mat Q0, const mat R0 );
	//! Here dt = [yt;ut] of appropriate dimensions
	void bayes ( const vec &dt );
	//! set estimates
	void set_est ( vec mu0, mat P0 ) {mu=mu0;P=P0;};
	//!dummy!
	const epdf& posterior() const{return E;};
	const enorm<fsqmat>* _e() const{return &E;};
	const mat _R() {return P;}
};

/*!
\brief Extended Kalman Filter

An approximation of the exact Bayesian filter with Gaussian noices and non-linear evolutions of their mean.
*/
template<class sq_T>
class EKF : public Kalman<fsqmat> {
	//! Internal Model f(x,u)
	diffbifn* pfxu;
	//! Observation Model h(x,u)
	diffbifn* phxu;
public:
	//! Default constructor
	EKF ( RV rvx, RV rvy, RV rvu );
	//! copy constructor
	EKF<sq_T>* _copy_() const { return new EKF<sq_T>(this); }
	//! Set nonlinear functions for mean values and covariance matrices.
	void set_parameters ( diffbifn* pfxu, diffbifn* phxu, const sq_T Q0, const sq_T R0 );
	//! Here dt = [yt;ut] of appropriate dimensions
	void bayes ( const vec &dt );
};

/*!
\brief Extended Kalman Filter in Square root

An approximation of the exact Bayesian filter with Gaussian noices and non-linear evolutions of their mean.
*/

class EKFCh : public KalmanCh {
protected:
	//! Internal Model f(x,u)
	diffbifn* pfxu;
	//! Observation Model h(x,u)
	diffbifn* phxu;
public:
	//! copy constructor duplicated - calls different set_parameters
	BM* _copy_() const {
		EKFCh* E=new EKFCh;
		E->set_parameters ( pfxu,phxu,Q,R );
		E->set_statistics ( _mu,_P );
		return E;
	}
	//! Set nonlinear functions for mean values and covariance matrices.
	void set_parameters ( diffbifn* pfxu, diffbifn* phxu, const chmat Q0, const chmat R0 );
	//! Here dt = [yt;ut] of appropriate dimensions
	void bayes ( const vec &dt );
};

/*!
\brief Kalman Filter with conditional diagonal matrices R and Q.
*/

class KFcondQR : public Kalman<ldmat> {
//protected:
public:
	void condition ( const vec &QR ) {
		it_assert_debug ( QR.length() == ( dimx+dimy ),"KFcondRQ: conditioning by incompatible vector" );

		Q.setD ( QR ( 0, dimx-1 ) );
		R.setD ( QR ( dimx, -1 ) );
	};
};

/*!
\brief Kalman Filter with conditional diagonal matrices R and Q.
*/

class KFcondR : public Kalman<ldmat> {
//protected:
public:
	//!Default constructor
	KFcondR ( ) : Kalman<ldmat> ( ) {};

	void condition ( const vec &R0 ) {
		it_assert_debug ( R0.length() == ( dimy ),"KFcondR: conditioning by incompatible vector" );

		R.setD ( R0 );
	};

};

//////// INstance

template<class sq_T>
Kalman<sq_T>::Kalman ( const Kalman<sq_T> &K0 ) : BM ( ),rvy ( K0.rvy ),rvu ( K0.rvu ),
		dimx ( K0.dimx ), dimy ( K0.dimy ),dimu ( K0.dimu ),
		A ( K0.A ), B ( K0.B ), C ( K0.C ), D ( K0.D ),
		Q ( K0.Q ), R ( K0.R ),
		est ( K0.est ), fy ( K0.fy ), _yp ( fy._mu() ),_Ry ( fy._R() ), _mu ( est._mu() ), _P ( est._R() ) {

// copy values in pointers
// 	_mu = K0._mu;
// 	_P = K0._P;
// 	_yp = K0._yp;
// 	_Ry = K0._Ry;

}

template<class sq_T>
Kalman<sq_T>::Kalman ( ) : BM (), est ( ), fy (),  _yp ( fy._mu() ), _Ry ( fy._R() ), _mu ( est._mu() ), _P ( est._R() ) {
};

template<class sq_T>
void Kalman<sq_T>::set_parameters ( const mat &A0,const  mat &B0, const mat &C0, const mat &D0, const sq_T &Q0, const sq_T &R0 ) {
	dimx = A0.rows();
	dimy = C0.rows();
	dimu = B0.cols();

	it_assert_debug ( A0.cols() ==dimx, "Kalman: A is not square" );
	it_assert_debug ( B0.rows() ==dimx, "Kalman: B is not compatible" );
	it_assert_debug ( C0.cols() ==dimx, "Kalman: C is not square" );
	it_assert_debug ( ( D0.rows() ==dimy ) || ( D0.cols() ==dimu ),	"Kalman: D is not compatible" );
	it_assert_debug ( ( R0.cols() ==dimy ) || ( R0.rows() ==dimy ), "Kalman: R is not compatible" );
	it_assert_debug ( ( Q0.cols() ==dimx ) || ( Q0.rows() ==dimx ), "Kalman: Q is not compatible" );

	A = A0;
	B = B0;
	C = C0;
	D = D0;
	R = R0;
	Q = Q0;
}

template<class sq_T>
void Kalman<sq_T>::bayes ( const vec &dt ) {
	it_assert_debug ( dt.length() == ( dimy+dimu ),"KalmanFull::bayes wrong size of dt" );

	sq_T iRy ( dimy );
	vec u = dt.get ( dimy,dimy+dimu-1 );
	vec y = dt.get ( 0,dimy-1 );
	//Time update
	_mu = A* _mu + B*u;
	//P  = A*P*A.transpose() + Q; in sq_T
	_P.mult_sym ( A );
	_P  +=Q;

	//Data update
	//_Ry = C*P*C.transpose() + R; in sq_T
	_P.mult_sym ( C, _Ry );
	_Ry  +=R;

	mat Pfull = _P.to_mat();

	_Ry.inv ( iRy ); // result is in _iRy;
	_K = Pfull*C.transpose() * ( iRy.to_mat() );

	sq_T pom ( ( int ) Pfull.rows() );
	iRy.mult_sym_t ( C*Pfull,pom );
	( _P ) -= pom; // P = P -PC'iRy*CP;
	( _yp ) = C* _mu  +D*u; //y prediction
	( _mu ) += _K* ( y- _yp );


	if ( evalll==true ) { //likelihood of observation y
		ll=fy.evallog ( y );
	}

//cout << "y: " << y-(*_yp) <<" R: " << _Ry->to_mat() << " iR: " << _iRy->to_mat() << " ll: " << ll <<endl;

};



//TODO why not const pointer??

template<class sq_T>
EKF<sq_T>::EKF ( RV rvx0, RV rvy0, RV rvu0 ) : Kalman<sq_T> ( rvx0,rvy0,rvu0 ) {}

template<class sq_T>
void EKF<sq_T>::set_parameters ( diffbifn* pfxu0,  diffbifn* phxu0,const sq_T Q0,const sq_T R0 ) {
	pfxu = pfxu0;
	phxu = phxu0;

	//initialize matrices A C, later, these will be only updated!
	pfxu->dfdx_cond ( _mu,zeros ( dimu ),A,true );
//	pfxu->dfdu_cond ( *_mu,zeros ( dimu ),B,true );
	B.clear();
	phxu->dfdx_cond ( _mu,zeros ( dimu ),C,true );
//	phxu->dfdu_cond ( *_mu,zeros ( dimu ),D,true );
	D.clear();

	R = R0;
	Q = Q0;
}

template<class sq_T>
void EKF<sq_T>::bayes ( const vec &dt ) {
	it_assert_debug ( dt.length() == ( dimy+dimu ),"KalmanFull::bayes wrong size of dt" );

	sq_T iRy ( dimy,dimy );
	vec u = dt.get ( dimy,dimy+dimu-1 );
	vec y = dt.get ( 0,dimy-1 );
	//Time update
	_mu = pfxu->eval ( _mu, u );
	pfxu->dfdx_cond ( _mu,u,A,false ); //update A by a derivative of fx

	//P  = A*P*A.transpose() + Q; in sq_T
	_P.mult_sym ( A );
	_P +=Q;

	//Data update
	phxu->dfdx_cond ( _mu,u,C,false ); //update C by a derivative hx
	//_Ry = C*P*C.transpose() + R; in sq_T
	_P.mult_sym ( C, _Ry );
	( _Ry ) +=R;

	mat Pfull = _P.to_mat();

	_Ry.inv ( iRy ); // result is in _iRy;
	_K = Pfull*C.transpose() * ( iRy.to_mat() );

	sq_T pom ( ( int ) Pfull.rows() );
	iRy.mult_sym_t ( C*Pfull,pom );
	( _P ) -= pom; // P = P -PC'iRy*CP;
	_yp = phxu->eval ( _mu,u ); //y prediction
	( _mu ) += _K* ( y-_yp );

	if ( evalll==true ) {ll+=fy.evallog ( y );}
};


}
#endif // KF_H