#include "libKF.h"

namespace bdm{

using std::endl;

KalmanFull::KalmanFull ( mat A0, mat B0, mat C0, mat D0, mat Q0, mat R0, mat P0, vec mu0 ) {
	dimx = A0.rows();
	dimu = B0.cols();
	dimy = C0.rows();

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

	A = A0;
	B = B0;
	C = C0;
	D = D0;
	R = R0;
	Q = Q0;
	mu = mu0;
	P = P0;
	
	ll = 0.0;
	evalll=true;
}

void KalmanFull::bayes ( const vec &dt ) {
	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;

	//Data update
	_Ry = C*P*C.transpose() + R;
	_iRy = inv ( _Ry );
	_K = P*C.transpose() *_iRy;
	P -= _K*C*P; // P = P -KCP;
	mu += _K* ( y-C*mu-D*u );

	if ( evalll ) {
		//from enorm
		vec x=y-C*mu-D*u;
		ll = -0.5* ( R.cols() * 1.83787706640935 +log ( det ( _Ry ) ) +x* ( _iRy*x ) );
	}
};

std::ostream &operator<< ( std::ostream &os, const KalmanFull &kf ) {
	os << "mu:" << kf.mu << endl << "P:" << kf.P <<endl;
	return os;
}



 /////////////////////////////// EKFS
EKFfull::EKFfull ( ) : BM (),E() {};

void EKFfull::set_parameters ( diffbifn* pfxu0,  diffbifn* phxu0,const mat Q0,const mat R0 ) {
	pfxu = pfxu0;
	phxu = phxu0;

	dimx = pfxu0->_dimx();
	dimy = phxu0->dimension();
	dimu = pfxu0->_dimu();
	E.epdf::set_parameters(dimx);
	
	A.set_size(dimx,dimx);
	C.set_size(dimy,dimx);
	//initialize matrices A C, later, these will be only updated!
	pfxu->dfdx_cond ( mu,zeros ( dimu ),A,true );
	B.clear();
	phxu->dfdx_cond ( mu,zeros ( dimu ),C,true );
	D.clear();

	R = R0;
	Q = Q0;
}

void EKFfull::bayes ( const vec &dt ) {
	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 );
	
	pfxu->dfdx_cond ( mu,zeros ( dimu ),A,true );
	phxu->dfdx_cond ( mu,zeros ( dimu ),C,true );

	//Time update
	mu = pfxu->eval(mu,u);// A*mu + B*u;
	P  = A*P*A.transpose() + Q;

	//Data update
	_Ry = C*P*C.transpose() + R;
	_iRy = inv ( _Ry );
	_K = P*C.transpose() *_iRy;
	P -= _K*C*P; // P = P -KCP;
	vec yp = phxu->eval(mu,u);
	mu += _K* ( y-yp );

	E.set_mu(mu);

	if ( BM::evalll ) {
		//from enorm
		vec x=y-yp;
		BM::ll = -0.5* ( R.cols() * 1.83787706640935 +log ( det ( _Ry ) ) +x* ( _iRy*x ) );
	}
};



void KalmanCh::set_parameters ( const mat &A0,const mat &B0,const mat &C0,const mat &D0,const chmat &Q0,const chmat &R0 ) {

	( ( Kalman<chmat>* ) this )->set_parameters ( A0,B0,C0,D0,Q0,R0 );
	// Cholesky special!
	preA=zeros(dimy+dimx+dimx,dimy+dimx);
//	preA.clear();
	preA.set_submatrix ( 0,0,R._Ch() );
	preA.set_submatrix ( dimy+dimx,dimy,Q._Ch() );
}


void KalmanCh::bayes ( const vec &dt ) {

	vec u = dt.get ( dimy,dimy+dimu-1 );
	vec y = dt.get ( 0,dimy-1 );
	vec pom(dimy);
	
	//TODO get rid of Q in qr()!
//	mat Q;

	//R and Q are already set in set_parameters()
	preA.set_submatrix ( dimy,0, ( _P._Ch() ) *C.T() );
	//Fixme can be more efficient if .T() is not used
	preA.set_submatrix ( dimy,dimy, ( _P._Ch() ) *A.T() );

//	if ( !qr ( preA,Q,postA ) ) it_warning ( "QR in kalman unstable!" );
	if ( !qr ( preA,postA ) ) {it_warning ( "QR in kalman unstable!" );}

	( _Ry._Ch() ) =postA ( 0,dimy-1, 0,dimy-1 );
	_K = inv(A)*( postA ( 0,dimy-1 ,dimy,dimy+dimx-1 ) ).T();
	( _P._Ch() ) = postA ( dimy,dimy+dimx-1,dimy,dimy+dimx-1 );
	
	_mu = A* ( _mu ) + B*u;
	_yp = C*_mu -D*u;
	
	backward_substitution(_Ry._Ch(),( y-_yp),pom);
	_mu +=  ( _K ) * pom;

/*		cout << "P:" <<_P.to_mat() <<endl;
		cout << "Ry:" <<_Ry.to_mat() <<endl;
		cout << "_K:" <<_K <<endl;*/
	
	if ( evalll==true ) { //likelihood of observation y
		ll=fy.evallog ( y );
	}
}



void EKFCh::set_parameters ( diffbifn* pfxu0,  diffbifn* phxu0,const chmat Q0,const chmat R0 ) {
	pfxu = pfxu0;
	phxu = phxu0;

	dimx = pfxu0->dimension();
	dimy = phxu0->dimension();
	dimu = pfxu0->_dimu();
	// set size of mu - just for constant terms of A and C 
	_mu=zeros(dimx);
	A=zeros(dimx,dimx);
	C=zeros(dimy,dimx);
	preA=zeros(dimy+dimx+dimx, dimy+dimx);
	
	//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;

	// Cholesky special!
	preA.clear();
	preA.set_submatrix ( 0,0,R._Ch() );
	preA.set_submatrix ( dimy+dimx,dimy,Q._Ch() );
}


void EKFCh::bayes ( const vec &dt ) {

	vec pom(dimy);
	vec u = dt.get ( dimy,dimy+dimu-1 );
	vec y = dt.get ( 0,dimy-1 );

	pfxu->dfdx_cond ( _mu,u,A,false ); //update A by a derivative of fx
	phxu->dfdx_cond ( _mu,u,C,false ); //update A by a derivative of fx

	//R and Q are already set in set_parameters()
	preA.set_submatrix ( dimy,0, ( _P._Ch() ) *C.T() );
	//Fixme can be more efficient if .T() is not used
	preA.set_submatrix ( dimy,dimy, ( _P._Ch() ) *A.T() );

// 	mat Sttm = _P->to_mat();
//	cout << preA <<endl;
//	cout << "_mu:" << _mu <<endl;

//	if ( !qr ( preA,Q,postA ) ) it_warning ( "QR in kalman unstable!" );
	if ( !qr ( preA,postA ) ) {it_warning ( "QR in kalman unstable!" );}


	( _Ry._Ch() ) =postA ( 0,dimy-1, 0,dimy-1 );
	_K = inv(A)*( postA ( 0,dimy-1 ,dimy,dimy+dimx-1 ) ).T();
	( _P._Ch() ) = postA ( dimy,dimy+dimx-1,dimy,dimy+dimx-1 );
	
// 	mat iRY = inv(_Ry->to_mat());
// 	mat Stt = Sttm - Sttm * C.T() * iRY * C * Sttm;
// 	mat _K2 = Stt*C.T()*inv(R.to_mat());

	// prediction
	_mu = pfxu->eval ( _mu ,u );
	_yp = phxu->eval ( _mu,u );
	
/*	vec mu2 = *_mu + ( _K2 ) * ( y-*_yp );*/
	
	//correction //= initial value is already prediction!
	backward_substitution(_Ry._Ch(),( y-_yp ),pom);
	_mu += ( _K ) *pom ;

/*	_K = (_P->to_mat())*C.transpose() * ( _iRy->to_mat() );
	*_mu = pfxu->eval ( *_mu ,u ) + ( _K )* ( y-*_yp );*/
	
// 		cout << "P:" <<_P.to_mat() <<endl;
// 		cout << "Ry:" <<_Ry.to_mat() <<endl;
// 	cout << "_mu:" <<_mu <<endl;
// 	cout << "dt:" <<dt <<endl;

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

}