#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> C1; C1.set_rv ( x );
	C1.set_parameters ( m1,V1 );
	enorm<fsqmat> C2; C2.set_rv ( x );
	C2.set_parameters ( m2,V2 );

	Array<epdf*> Sim ( 2 ); Sim ( 0 ) =&C1;Sim ( 1 ) =&C2;
	emix Simul;Simul.set_rv ( x );
	Simul.set_parameters ( "0.5 0.6", Sim, false );

	// 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;
}