#include "kalman.h"

namespace bdm {

using std::endl;



void KalmanFull::bayes ( const vec &yt, const vec &cond ) {
	bdm_assert_debug ( yt.length() == ( dimy ), "KalmanFull::bayes wrong size of dt" );

	const vec &u = cond; // in this case cond=ut
	const vec &y = yt;
	
	vec& mu = est._mu();
	mat &P = est._R();
	mat& _Ry = fy._R();
	vec& yp = fy._mu();
	//Time update
	mu = A * mu + B * u;
	P  = A * P * A.transpose() + (mat)Q;

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

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



/////////////////////////////// EKFS
EKFfull::EKFfull ( ) : KalmanFull () {};

void EKFfull::set_parameters ( const shared_ptr<diffbifn> &pfxu0, const shared_ptr<diffbifn> &phxu0, const mat Q0, const mat R0 ) {
	pfxu = pfxu0;
	phxu = phxu0;

	set_dim( pfxu0->_dimx());
	dimy = phxu0->dimension();
	dimc = pfxu0->_dimu();
	est.set_parameters( zeros(dimension()), eye(dimension()) );

	A.set_size ( dimension(), dimension() );
	C.set_size ( dimy, dimension() );
	//initialize matrices A C, later, these will be only updated!
	pfxu->dfdx_cond ( est._mu(), zeros ( dimc ), A, true );
	B.clear();
	phxu->dfdx_cond ( est._mu(), zeros ( dimc ), C, true );
	D.clear();

	R = R0;
	Q = Q0;
}

void EKFfull::bayes ( const vec &yt, const vec &cond) {
	bdm_assert_debug ( yt.length() == ( dimy ), "EKFull::bayes wrong size of dt" );
	bdm_assert_debug ( cond.length() == ( dimc ), "EKFull::bayes wrong size of dt" );
	
	const vec &u = cond;
	const vec &y = yt; //lazy to change it
	vec &mu = est._mu();
	mat &P = est._R();
	mat& _Ry = fy._R();
	vec& yp = fy._mu();
	
	pfxu->dfdx_cond ( mu, zeros ( dimc ), A, true );
	phxu->dfdx_cond ( mu, zeros ( dimc ), C, true );

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

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

	if ( BM::evalll ) {
		ll=fy.evallog(y);
	}
};



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

	( ( StateSpace<chmat>* ) this )->set_parameters ( A0, B0, C0, D0, Q0, R0 );
	
	_K=zeros(dimension(),dimy);
}

void KalmanCh::initialize(){
	preA = zeros ( dimy + dimension() + dimension(), dimy + dimension() );
//	preA.clear();
	preA.set_submatrix ( 0, 0, R._Ch() );
	preA.set_submatrix ( dimy + dimension(), dimy, Q._Ch() );
}

void KalmanCh::bayes ( const vec &yt, const vec &cond ) {
	bdm_assert_debug(yt.length()==dimy, "yt mismatch");
	bdm_assert_debug(cond.length()==dimc, "yt mismatch");
	
	const vec &u = cond;
	const vec &y = yt;
	vec pom ( dimy );

	chmat &_P=est._R();
	vec &_mu = est._mu();
	mat _K(dimension(),dimy);
	chmat &_Ry=fy._R();
	vec &_yp = fy._mu();
	//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, postA ) ) {
		bdm_warning ( "QR in KalmanCh unstable!" );
	}

	( _Ry._Ch() ) = postA ( 0, dimy - 1, 0, dimy - 1 );
	_K = inv ( A ) * ( postA ( 0, dimy - 1 , dimy, dimy + dimension() - 1 ) ).T();
	( _P._Ch() ) = postA ( dimy, dimy + dimension() - 1, dimy, dimy + dimension() - 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 ( const shared_ptr<diffbifn> &pfxu0, const shared_ptr<diffbifn> &phxu0, const chmat Q0, const chmat R0 ) {
	pfxu = pfxu0;
	phxu = phxu0;

	set_dim( pfxu0->_dimx());
	dimy = phxu0->dimension();
	dimc = pfxu0->_dimu();
	
	vec &_mu = est._mu();
	// if mu is not set, set it to zeros, just for constant terms of A and C
	if ( _mu.length() != dimension() ) _mu = zeros ( dimension() ); 
	A = zeros ( dimension(), dimension() );
	C = zeros ( dimy, dimension() );
	preA = zeros ( dimy + dimension() + dimension(), dimy + dimension() );

	//initialize matrices A C, later, these will be only updated!
	pfxu->dfdx_cond ( _mu, zeros ( dimc ), A, true );
//	pfxu->dfdu_cond ( *_mu,zeros ( dimu ),B,true );
	B.clear();
	phxu->dfdx_cond ( _mu, zeros ( dimc ), 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 + dimension(), dimy, Q._Ch() );
}


void EKFCh::bayes ( const vec &yt, const vec &cond ) {

	vec pom ( dimy );
	const vec &u = cond;
	const vec &y = yt;
	vec &_mu = est._mu();
	chmat &_P = est._R();
	chmat &_Ry = fy._R();
	vec &_yp = fy._mu();
	
	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, postA ) ) {
		bdm_warning ( "QR in EKFCh unstable!" );
	}


	( _Ry._Ch() ) = postA ( 0, dimy - 1, 0, dimy - 1 );
	_K = inv ( A ) * ( postA ( 0, dimy - 1 , dimy, dimy + dimension() - 1 ) ).T();
	( _P._Ch() ) = postA ( dimy, dimy + dimension() - 1, dimy, dimy + dimension() - 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 );
	}
}

void EKFCh::from_setting ( const Setting &set ) {
	shared_ptr<diffbifn> IM = UI::build<diffbifn> ( set, "IM", UI::compulsory );
	shared_ptr<diffbifn> OM = UI::build<diffbifn> ( set, "OM", UI::compulsory );

	//statistics
	int dim = IM->dimension();
	vec mu0;
	if ( !UI::get ( mu0, set, "mu0" ) )
		mu0 = zeros ( dim );

	mat P0;
	vec dP0;
	if ( UI::get ( dP0, set, "dP0" ) )
		P0 = diag ( dP0 );
	else if ( !UI::get ( P0, set, "P0" ) )
		P0 = eye ( dim );

	set_statistics ( mu0, P0 );

	//parameters
	vec dQ, dR;
	UI::get ( dQ, set, "dQ", UI::compulsory );
	UI::get ( dR, set, "dR", UI::compulsory );
	set_parameters ( IM, OM, diag ( dQ ), diag ( dR ) );

	//connect
	shared_ptr<RV> drv = UI::build<RV> ( set, "drv", UI::compulsory );
	set_yrv ( *drv );
	shared_ptr<RV> rv = UI::build<RV> ( set, "rv", UI::compulsory );
	set_rv ( *rv );

	string options;
	if ( UI::get ( options, set, "options" ) )
		set_options ( options );
}

void MultiModel::from_setting ( const Setting &set ) {
	Array<EKFCh*> A;
	UI::get ( A, set, "models", UI::compulsory );

	set_parameters ( A );
	set_yrv ( A ( 0 )->_yrv() );
	//set_rv(A(0)->_rv());

	string options;
	if ( set.lookupValue ( "options", options ) )
		set_options ( options );
}

}