
#include "kalman.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();
	// if mu is not set, set it to zeros, just for constant terms of A and C
	if ( _mu.length() != dimx ) _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 );
	}
}

void EKFCh::from_setting ( const Setting &set ) {
	diffbifn* IM = UI::build<diffbifn> ( set, "IM", UI::compulsory );
	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
	RV* drv = UI::build<RV> ( set, "drv", UI::compulsory );
	set_drv ( *drv );
	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_drv ( A ( 0 )->_drv() );
	//set_rv(A(0)->_rv());

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

}
