#include "estim/arx.h"
#include "../mat_checks.h"

using namespace bdm;

TEST ( arx_elem_stress ) {
	// Setup model : ARX for 1D Gaussian
	//Test constructor
	mat V0 = 0.00001 * eye ( 2 );
	V0 ( 0, 0 ) = 0.1; //
	ARX Ar;
	Ar.set_statistics ( 1, V0, -1.0 );
	Ar.set_constant ( true );
	Ar.validate();

	mat mu ( 1, 1 );
	mat R ( 1, 1 );
	Ar.posterior().mean_mat ( mu, R );
	cout << "Prior moments: mu=" << mu << ", R=" << R << endl;

	int ndat = 200;
	vec smp = randn ( ndat );
	//
	mat Smp = ones ( 2, ndat );
	Smp.set_row ( 0, smp );
	//
	Ar.bayes_batch ( Smp );
	// Ar is now filled with estimates of N(0,1);
	cout << "Empirical moments: mu=" << sum ( smp ) / ndat << ", R=" << sum_sqr ( smp ) / ndat - pow ( sum ( smp ) / ndat, 2 ) << endl;
	Ar.posterior().mean_mat ( mu, R );
	cout << "Posterior moments: mu=" << mu << ", R=" << R << endl;

	//////// TEST prediction
	vec x = linspace ( -3.0, 3.0, 100 );
	double xstep = 6.0 / 100.0;
	mat X ( 1, 100 );
	mat X2 ( 2, 100 );
	X.set_row ( 0, x );
	X2.set_row ( 0, x );

	mlstudent* Ap = Ar.predictor_student();
	vec Ap_x = Ap->evallogcond_mat ( X, empty_vec );
	vec ll_x = Ar.logpred_mat ( X2 );

	cout << "normalize : " << xstep*sum ( exp ( Ap_x ) ) << endl;
	cout << "normalize : " << xstep*sum ( exp ( ll_x ) ) << endl;

	it_file it ( "arx_elem_test.it" );
	it << Name ( "Ap_x" ) << Ap_x;
	it << Name ( "ll_x" ) << ll_x;
}
