/*!
  \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 <itpp/itbase.h>
#include "libBM.h"
#include "libDC.h"


using namespace itpp;

/*!
* \brief Basic Kalman filter with full matrices (education purpose only)!
*/
class KalmanFull : public BM { 
	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;

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, bool evalll=true); 

	friend std::ostream &operator<< ( std::ostream &os, const KalmanFull &kf );

};


/*!
* \brief Kalman filter with covaraince matrices in square root form.
*/
template<class sq_T>
class Kalman : public BM { 
	int dimx, dimy, dimu;
	mat A, B, C, D;
	sq_T R, Q;
	
	//cache
	mat _K, _yp;
	sq_T _Ry,_iRy;
public:
	//posterior 
	//! Mean value of the posterior density
	vec mu;
	//! Mean value of the posterior density
	sq_T P;

public:
	//! Full constructor 
	Kalman ( mat A0, mat B0, mat C0, mat D0, sq_T R0, sq_T Q0, sq_T P0, vec mu0 );
	//! Here dt = [yt;ut] of appropriate dimensions
	void bayes(const vec &dt, bool evalll=true); 

	friend std::ostream &operator<< ( std::ostream &os, const KalmanFull &kf );

};

//////// INstance

template<class sq_T>
Kalman<sq_T>::Kalman( mat A0, mat B0, mat C0, mat D0, sq_T R0, sq_T Q0, sq_T P0, vec mu0 ) {
	dimx = A0.rows();
	dimu = B0.cols();
	dimy = C0.rows();

	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;
	mu = mu0;
	P = P0;

	ll = 0;
//Fixme should we assign cache??
}

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

	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_qform( A );
	P+=Q;

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

	mat Pfull = P.to_mat();
	
	_Ry.inv( _iRy ); // result is in _iRy;
	_K = Pfull*C.transpose()*(_iRy.to_mat());
	P -= _K*C*Pfull; // P = P -KCP;
	_yp = y-C*mu-D*u; //y prediction
	mu += _K*( _yp );
	
	if (evalll==true) {
	ll+= -0.5*_Ry.logdet() -0.5*_iRy.qform(_yp);
	}
};

//extern template class Kalman<ldmat>; 


#endif // KF_H