#include <estim/arx.h>
#include <estim/merger.h>
#include <stat/libEF.h>
#include <stat/loggers.h>
//#include <stat/merger.h>
using namespace itpp;

//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.count() ); V0 ( 0,0 ) *= 10; //
	mat V0g = 0.001*eye ( thrg.count() ); V0g ( 0,0 ) *= 10; //
	double nu0 = ord+6.0;
	double frg = 1.0;//0.99;

	ARX P1 ( thri, V0, nu0, frg );
	ARX P2 ( thri, V0, nu0, frg );
	ARX PG ( thrg, V0g, nu0, frg );
	//Test estimation
	int ndat = 500;
	int t;

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

	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 ln}" ), "Pred" );

	L.init();

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

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

	vec PredLLs(5);
	vec PostLLs(3);
	vec PredLLs_m=zeros(5);
	vec PostLLs_m=zeros(3);
	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/40.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_student(y,concat(ym,u1));
				mlnorm<ldmat>* P2p = P2.predictor_student(y,concat(ym,u2));
				mlnorm<ldmat>* Pgp = PG.predictor_student(y,concat(ym,uu));
				
				Array<mpdf*> A(2); A(0)=P1p;A(1)=P2p;
				merger M(A);
				enorm<ldmat> g0(concat(yy,uu)); g0.set_parameters("0 0 0 0 ",3*eye(4));
				M.set_parameters(1.2, 100,1);
				M.merge(&g0);
				
				yt(0) = Yt(t);
				double P1pl = P1p->evalcond(yt,rgr1);	
				double P2pl = P2p->evalcond(yt,rgr2);	
				double PGpl = Pgp->evalcond(yt,rgrg);	
				PredLLs(0) = log(P1pl); PredLLs(1) = log(P2pl); PredLLs(2) = log(PGpl);
				{
					ARX* Apred = (ARX*)M._Mix()._Coms(0);
					enorm<ldmat>* MP= Apred->predictor(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);

					double cP1pl = cP1p->evalcond(yt,rgr1);	
					double cP2pl = cP2p->evalcond(yt,rgr2);	
					PredLLs(3) = log(cP1pl);
					PredLLs(4) = log(cP2pl);
				}

				L.logit(Li_Pred, PredLLs_m*frg+ PredLLs); //log-normal
				PredLLs_m *=frg;
				PredLLs_m += PredLLs;
				
				delete P1p;
				delete P2p;
				delete Pgp;
			}
			
			// 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._epdf().mean() );
		L.logit ( Li_Eth2, P2._epdf().mean() );
		L.logit ( Li_Ethg, PG._epdf().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" );
}