#include "estim/mixtures.h"
#include "estim/arx.h"
#include "stat/exp_family.h"
using namespace bdm;

//These lines are needed for use of cout and endl
using std::cout;
using std::endl;

void disp_mix2D ( const MixEF &Mi ) {
	const eprod &ep = ( ( const eprod& ) ( Mi.posterior() ) );
	int i;
	mat Var ( 2, 2 ), M ( 1, 2 );
	for ( i = 0; i < ep._rv().length() - 1; i++ ) { // -1 => last one is the weight
		( ( egiw* ) ep ( i ) )->mean_mat ( M, Var );
		cout << "Mean: " << M << endl << "Variance: " << endl << Var << endl;
	}
	cout << "Weights: " << ep ( i )->mean() << endl;
}

mat grid2D ( const MixEF &M, const vec &xb, const vec &yb, int Ngr = 100 ) {
	mat PPdf ( Ngr + 1, Ngr + 1 );
	vec rgr ( 3 );

	rgr ( 2 ) = 1;
	int i = 0, j = 0;
	double xstep = ( xb ( 1 ) - xb ( 0 ) ) / Ngr;
	double ystep = ( yb ( 1 ) - yb ( 0 ) ) / Ngr;

	for ( double x = xb ( 0 ); x <= xb ( 1 ); x += xstep, i++ ) {
		rgr ( 0 ) = x;
		j = 0;
		for ( double y = yb ( 0 ); y <= yb ( 1 ); y += ystep, j++ ) {
			rgr ( 1 ) = y;
			PPdf ( i, j ) = exp ( M.logpred ( rgr ) );
		}
	}
	return PPdf;
}

int main() {
	RV x ( "{x }", "2" );


	cout << "Test for estimation of a 2D Gaussian mixture" << endl;
	///////////////////////////////////
	// Setup model

	vec m1 ( "-2 -2" );
	vec m2 ( "2 2" );

	fsqmat V1 ( mat ( "1 0.5; 0.5 1" ) );
	fsqmat V2 ( mat ( "2 -0.1; -0.1 2" ) );

	enorm_fsqmat_ptr C1;
	C1->set_rv ( x );
	C1->set_parameters ( m1, V1 );
	enorm_fsqmat_ptr C2;
	C2->set_rv ( x );
	C2->set_parameters ( m2, V2 );

	epdf_array Sim ( 2 );
	Sim ( 0 ) = C1;
	Sim ( 1 ) = C2;
	emix Simul;
	Simul.set_rv ( x );
	Simul.set_parameters ( "0.5 0.6", Sim );

	// Sample parameters
	int ndat = 100;
	mat Smp = Simul.sample_m ( ndat );

	cout << "Simulated means: " << m1 << " and " << m2 << endl;
	cout << "Simulated covariances: " << endl << V1 << " and " << endl << V2 << endl;

	//////////////////////////////
	// Prepare estimator manually

	// Initial values of components
	mat V0 = 1e-4 * eye ( 3 );
	V0 ( 0, 0 ) *= 100;
	V0 ( 1, 1 ) *= 100;
	mat Vp = "0 0 1; 0 0 1; 1 1 0";
	Vp *= 5 * 1e-5;

	ARX M1;
	M1.set_statistics ( 2, V0 - Vp, 8 );
	ARX M2;
	M2.set_statistics ( 2, V0 + Vp, 8 );

	// Build mixture model
	Array<BMEF*> A ( 2 );
	A ( 0 ) = &M1;
	A ( 1 ) = &M2;
	MixEF Post ( A, "0.5 0.5" );
	cout << "Initial mixture:" << endl;
	disp_mix2D ( Post );

	// Add ones for constant coefficients
	mat Data = concat_vertical ( Smp, ones ( 1, Smp.cols() ) );
	Post.bayes ( Data );

	cout << "Posterior mixture:" << endl;
	disp_mix2D ( Post );

	//TEST random initialization
	RNG_randomize();
	ARX Aflat;
	Aflat.set_statistics ( 2, V0, 7 );
	MixEF RND;
	RND.init ( &Aflat, Data, 10 ); // already initialized!
	cout << endl << "== Randomly initialized mixture ==" << endl;
	cout << endl << "== INIT ==" << endl;
	disp_mix2D ( RND );
	mat PPdf_I = grid2D ( RND, vec_2 ( -5.0, 5.0 ), vec_2 ( -5.0, 5.0 ) );;

	RND.bayes ( Data );
	cout << endl << "== BAYES ==" << endl;
	disp_mix2D ( RND );


	it_file of ( "mixef_test.it", true );
	of << Name ( "Smp" ) << Smp;
	of << Name ( "PPdf" ) << grid2D ( Post, vec_2 ( -5.0, 5.0 ), vec_2 ( -5.0, 5.0 ) );
	of << Name ( "PPdf_I" ) << PPdf_I;
	of << Name ( "PPdf_RND" ) << grid2D ( RND, vec_2 ( -5.0, 5.0 ), vec_2 ( -5.0, 5.0 ) );

	cout << endl << "Variables written to file mixef_test.it:" << endl;
	cout << "Smp  : data sampled from the simulator " << endl;
	cout << "PPdf : posterior density on a grid <-5,5><-5,5>" << endl;
	cout << "PPdf_RND : posterior density of randomly initilaized mixture on a grid <-5,5><-5,5>" << endl;
}