/*!
  \file
  \brief Disturbances of PMSM model are fitted by an ARX model
  \author Vaclav Smidl.

  \ingroup PMSM
  -----------------------------------
  BDM++ - C++ library for Bayesian Decision Making under Uncertainty

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


#include <estim/arx.h>

#include <stat/loggers.h>

using namespace bdm;

vec getPsi_a ( int t, mat &D, mat &Du , mat &X );
vec getPsi_b ( int t, mat &D, mat &Du , mat &X );

int main() {
	// Kalman filter
	int Ndat = 90000;
	mat D;
	mat Du;
	mat X;

	dirfilelog L ( "exp/sim_var_arx",1000 );

	it_file itf ( "sim_var.it" );
	itf >> Name ( "D" ) >> D;
	itf >> Name ( "Du" ) >> Du;
	itf >> Name ( "X" ) >> X;

	Array<std::string> Names = "{ia ib om dom r }";
	int rglen = Names.length();
	//Regressor
	RV rgr ( Names );
	mat V0 = 0.0001*eye ( rglen ); V0 ( 0,0 ) =200;
	double nu0 = rglen+1;

	//Autoregressive model
	ARX Ar_a ( rgr,V0,nu0 ,0.95 );
	ARX Ar_b ( rgr,V0,nu0 ,0.95 );
	const epdf& pA= Ar_a.posterior();
	const epdf& pB= Ar_b.posterior();

	RV rta ( "{th_a }",vec_1 ( rglen ) );
	RV rtb ( "{th_b }",vec_1 ( rglen ) );
	int tha_log = L.add ( rta,"" );
	int thb_log = L.add ( rtb,"" );
	L.init();

	vec Psi ( rglen );
	vec Save ( 13 );
	for ( int t=2; t<Ndat; t++ ) {
		Psi =getPsi_a ( t, D,Du ,X );
		Ar_a.bayes ( Psi );
		Psi =getPsi_b ( t, D,Du ,X );
		Ar_b.bayes ( Psi );

		Save = pA.mean();
		L.logit ( tha_log, Save );
		Save = pB.mean();
		L.logit ( thb_log, Save );
		L.step();
	}
	L.finalize();

	ivec bestind = Ar_a.structure_est ( egiw ( rgr,V0,nu0 ) );
	ivec bestind2 = Ar_b.structure_est ( egiw ( rgr,V0,nu0 ) );
	cout << bestind <<endl;
	cout << bestind2 <<endl;

	return 0;
}

vec getPsi_a ( int t, mat &D, mat &Du, mat &X ) {
	vec Psi ( 5 );
	Psi ( 0 ) = D ( t,0 )-Du ( t,0 ); // a=0

	Psi ( 1 ) = D ( t,2 );
	Psi ( 2 ) = D(t,2)-D(t-1,2 );
	Psi ( 3 ) = D(t,0)-D(t-1,0 );
	Psi ( 4 ) = X ( t,2 ) - X ( t-1,2 );

	return Psi;
}

vec getPsi_b ( int t, mat &D, mat &Du, mat &X ) {
	vec Psi ( 5 );
	Psi ( 0 ) = D ( t,1 )-Du ( t,1 ); //b=1

	Psi ( 1 ) = D(t,3);
	Psi ( 2 ) = D(t,3)-D(t-1,3 );
	Psi ( 3 ) = D ( t,1 )-D(t-1,1);
	Psi ( 4 ) = X ( t,2 ) - X ( t-1,2 );

	return Psi;
}