#include <estim/arx.h>
#include <estim/merger.h>
#include <stat/exp_family.h>
#include <stat/loggers.h>
//#include <stat/merger.h>
using namespace bdm;

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

int main() {
	// Setup model
	RV y ( "{y }" );
	RV u1 ( "{u1 }" );
	RV u2 ( "{u2 }" );
	RV uu=u1; uu.add ( u2 );

	double a1t = 1.5;
	double a2t = 0.8;
	double sqr=0.10;
	// Full system
	vec thg =vec_2 ( a1t,a2t ); //Simulated system - zero for constant term
	vec Th = concat ( thg, sqr ); //Full parameter

	// Estimated systems ARX(2)
	RV a1 ( "{a1 }" );
	RV a2 ( "{a2 }" );
	RV r ( "{r }" );
	RV all =a1; all.add ( a2 ); all.add ( r );
	RV allj =a1; allj.add ( r ); allj.add ( a2 );
	vec Thj=vec_3 ( a1t,sqr,a2t );
	// Setup values

	//ARX constructor
	mat V0 = 0.001*eye ( 2 ); V0 ( 0,0 ) = 1; //
	mat V0g = 0.001*eye ( 3 ); V0g ( 0,0 ) = 1; //

	ARX P1; P1.set_statistics(2, V0, -1 );
	ARX P2; P2.set_statistics(2, V0, -1 );
	ARX PG; PG.set_statistics(3, V0g, -1 ); //or -1?
//	ARX PGk ( all, V0g, -1 );

	//Test estimation
	int ndat = 100;
	int t;

	// Logging
	dirfilelog L ( "exp/merg_giw",ndat );
	int Li_Data = L.add ( RV ( "{Y U1 U2 }" ), "" );
	int Li_LL   = L.add ( RV ( "{G M }" ), "LL" );
	int Li_P1m   = L.add ( RV ( "{a1 r }" ), "P1" );
	int Li_P2m   = L.add ( RV ( "{a2 r }" ), "P2" );
	int Li_Gm   = L.add ( RV ( "{a1 a2 r }" ), "G" );
	int Li_Mm   = L.add ( RV ( "{a1 r a2 }" ), "M" );
	L.init();

	vec Yt ( ndat );
	vec yt ( 1 );

	vec LLs ( 2 );
	vec rgrg ( 2 );

	//Proposal
	enorm<ldmat> g0; g0.set_rv( a1 ); g0.set_parameters ( "1  ",mat ( "1" ) );
	egamma g1; g1.set_rv ( r ); g1.set_parameters ( "2  ", "2" );
	enorm<ldmat> g2; g2.set_rv ( a2 ); g2.set_parameters ( "1  ",mat ( "1" ) );

	Array<const epdf*> A ( 3 ); A ( 0 ) = &g0; A ( 1 ) =&g1; A ( 2 ) = &g2;
	eprod G0; G0.set_parameters ( A );

	for ( t=0; t<ndat; t++ ) {
		// True system
		rgrg ( 0 ) = pow ( sin ( ( t/40.0 ) *pi ),3 );
		rgrg ( 1 ) = pow ( cos ( ( t/40.0 +0.1 ) *pi ),3 );

		Yt ( t ) = thg*rgrg + sqr * NorRNG();

		// Bayes for all
		P1.bayes ( concat ( Yt ( t ),vec_1 ( rgrg ( 0 ) ) ) );
		P2.bayes ( concat ( Yt ( t ),vec_1 ( rgrg ( 1 ) ) ) );
		PG.bayes ( concat ( Yt ( t ),rgrg ) );


		// crippling PGk by substituting zeros.
		/*	ldmat &Vk=const_cast<egiw*>(PGk._e())->_V();  //PG ldmat does not like 0!
			mat fVk=PG._e()->_V().to_mat();
			fVk(1,2) = 0.0;
			fVk(2,1) = 0.0;
			Vk = ldmat(fVk);
		*/	//PGk is now krippled

		// Merge estimates
		Array<pdf*> A ( 2 ); A ( 0 ) =&P1;A ( 1 ) =&P2;
		merger M ( A );
		M.set_parameters ( 1.5, 100,3 ); //M._Mix().set_method(QB);
		//M2.set_parameters ( 100.0, 1000,3 ); //M2._Mix().set_method(QB);
		M.merge ( &G0 );
		//M2.merge ( &G0 );

		//Likelihood
		yt ( 0 ) = Yt ( t );

//		LLs ( 0 ) = P1._e()->evallog ( get_vec(Th, "1 2") );
//		LLs ( 1 ) = P2._e()->evallog ( get_vec(Th, "3 2")  );
		LLs ( 0 ) = PG._e()->evallog ( Th );
		LLs ( 1 ) = M._Mix().logpred ( concat ( Thj, vec_1 ( 1.0 ) ) );
//		LLs ( 3 ) = M2._Mix().logpred ( concat(Thj, vec_1(1.0)) );
		L.logit ( Li_LL, LLs ); //log-normal

		//Logger
		L.logit ( Li_Data, vec_3 ( Yt ( t ), rgrg ( 0 ), rgrg ( 1 ) ) );
		L.logit ( Li_P1m, P1._e()->mean() );
		L.logit ( Li_P2m, P2._e()->mean() );
		L.logit ( Li_Gm, PG._e()->mean() );
		L.logit ( Li_Mm, M.mean() );

		L.step ( );

		cout << "Vg: " << PG._e()->_V().to_mat() <<endl;
		vec mea = M.mean();
		cout << "Ve: " << M.variance() <<endl;
	}
	L.finalize( );
	L.itsave ( "merg_egiw.it" );
	cout << endl;
}