#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 ym=y; ym.t ( -1 );
	RV yy = y; yy.add ( ym );
	RV uu=u1; uu.add ( u2 );

	// Full system
	vec thg ( "0.7 1 1 0" ); //Simulated system - zero for constant term
	//y=a y_t-1 + u1 + u2
	double sqr=0.1;
	int ord = 1;

	// Estimated systems ARX(2)
	RV thri ( "{thr_i }",vec_1 ( 3+1 ) );
	RV thrg ( "{thr_g }",vec_1 ( 4+1 ) );
	// Setup values

	//ARX constructor
	mat V0 = 0.001*eye ( thri._dsize() ); V0 ( 0,0 ) *= 10; //
	mat V0g = 0.001*eye ( thrg._dsize() ); V0g ( 0,0 ) *= 10; //
	double nu0 = ord+6.0;
	double frg = 0.95;

	ARX P1; P1.set_statistics( 2, V0, nu0); P1.set_parameters(frg );
	ARX P2;P1.set_statistics(2, V0, nu0); P2.set_parameters( frg );
	ARX PG;PG.set_statistics(3, V0g, nu0);; PG.set_parameters( frg );
	//Test estimation
	int ndat = 200;
	int t;

	// Logging
	dirfilelog L ( "exp/merg",10 );

	int Li_Eth1 = L.add ( thri,"P1" );
	int Li_Eth2 = L.add ( thri,"P2" );
	int Li_Ethg = L.add ( thrg,"PG" );
	int Li_Data = L.add ( RV ( "{Y U1 U2 }" ), "" );
	int Li_LL   = L.add ( RV ( "{1 2 G }" ), "LL" );
	int Li_Pred = L.add ( RV ( "{1 2 G ar ge }" ), "Pred" );

	L.init();

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

	Yt.set_subvector ( 0,randn ( ord ) ); //initial values
	vec rgrg ( thrg._dsize() -1 ); // constant terms are in!
	vec rgr1 ( thri._dsize() -1 );
	vec rgr2 ( thri._dsize() -1 );

	vec PredLLs ( 5 );
	vec PostLLs ( 3 );
	vec PredLLs_m=zeros ( 5 );
	ivec ind_r1 = "0 1 3";
	ivec ind_r2 = "0 2 3";
	for ( t=0; t<ndat; t++ ) {
		// True system
		if ( t>0 ) {
			rgrg ( 0 ) =Yt ( t-1 );
			rgrg ( 1 ) = pow(sin ( ( t/40.0 ) *pi ),3);
			rgrg ( 2 ) = pow(cos ( ( t/60.0 ) *pi ),3);
			rgrg ( 3 ) = 1.0; // constant term

			rgr1 ( 0 ) = rgrg ( 0 ); rgr1 ( 1 ) = rgrg ( 1 ); rgr1 ( 2 ) = rgrg ( 3 ); // no u2
			rgr2 ( 0 ) = rgrg ( 0 ); rgr2 ( 1 ) = rgrg ( 2 ); rgr2 ( 2 ) = rgrg ( 3 ); // no u1

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

			// Test predictors
			if ( t>2 ) {
					mlnorm<ldmat>* P1p = P1.predictor();// ( y,concat ( ym,u1 ) );
					mlnorm<ldmat>* P2p = P2.predictor();// ( y,concat ( ym,u2 ) );
					mlnorm<ldmat>* Pgp = PG.predictor();// ( y,concat ( ym,uu ) );

					Array<pdf*> A ( 2 ); A ( 0 ) =P1p;A ( 1 ) =P2p;
					merger M ( A );
					enorm<ldmat> g0;g0.set_rv ( concat ( yy,uu ) ); g0.set_parameters ( "0 0 0 0 ",3*eye ( 4 ) );
					M.set_parameters ( 1e8, 101,1 );
					M.merge ( &g0 );
				yt ( 0 ) = Yt ( t );
				double P1pl = P1p->evallogcond ( yt,rgr1 );
				double P2pl = P2p->evallogcond ( yt,rgr2 );
				double PGpl = Pgp->evallogcond ( yt,rgrg );
				{
					cout << "T: " << t << "yt: " << yt << endl;
					cout << "yt_1: " << P1p->_epdf().mean() << endl;
					cout << *P1p <<endl;
					cout << "yt_2: " << P2p->_epdf().mean() << endl;
					cout << "yt_G: " << P2p->_epdf().mean() << endl;
				}
				double cP1pl;
				double cP2pl;
				{
					ARX* Apred = ( ARX* ) M._Mix()._Coms ( 0 );
					enorm<ldmat>* MP= Apred->epredictor ();//( concat ( yy,uu ) );
					enorm<ldmat>* mP1p = ( enorm<ldmat>* ) MP->marginal ( concat ( yy,u1 ) );
					enorm<ldmat>* mP2p = ( enorm<ldmat>* ) MP->marginal ( concat ( yy,u2 ) );
					mlnorm<ldmat>* cP1p = ( mlnorm<ldmat>* ) mP1p->condition ( y );
					mlnorm<ldmat>* cP2p = ( mlnorm<ldmat>* ) mP2p->condition ( y );

					cP1pl = cP1p->evallogcond ( yt,rgr1 );
					cP2pl = cP2p->evallogcond ( yt,rgr2 );

					cout << "ytm1: " << cP1p->_epdf().mean() << endl;
					cout << "ytm2: " << cP2p->_epdf().mean() << endl;
					cout << *cP2p << endl;
				}

				PredLLs *=frg;
				PredLLs += concat ( vec_3 ( P1pl, P2pl, PGpl ), vec_2 ( cP1pl,cP2pl ) );
				L.logit ( Li_Pred, PredLLs ); //log-normal

				delete P1p;
				delete P2p;
				delete Pgp;
			}

			if ( t<100 ) {
				// 1st
				P1.bayes ( concat ( Yt ( t ),rgr1 ) );
				// 2nd
				P2.bayes ( concat ( Yt ( t ),rgr2 ) );

				//Global
				PG.bayes ( concat ( Yt ( t ),rgrg ) );
			}
			//Merger
		}
		L.logit ( Li_Eth1, P1.posterior().mean() );
		L.logit ( Li_Eth2, P2.posterior().mean() );
		L.logit ( Li_Ethg, PG.posterior().mean() );
		L.logit ( Li_Data, vec_3 ( Yt ( t ), rgrg ( 1 ), rgrg ( 2 ) ) );
		PostLLs *= frg;
		PostLLs += vec_3 ( P1._ll(), P2._ll(), PG._ll() );
		L.logit ( Li_LL, PostLLs );
		L.step ( );
	}
	L.finalize( );
	L.itsave ( "merg.it" );
}