#include <estim/arx.h>
#include <estim/merger.h>
#include <stat/libEF.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;

/*! The purpose of this experiment is to test merging fragmental pdfs between two participants sharing the same parameter, a. However, this parameter has a different role in each model.

Lets assume that:
 u -> |a,r| -> y -> |a,r| -> z
*/

int main() {
	// Setup model
	RV y ( "{y }" );
	RV u ( "{u }" );
	RV z ( "{z }" );
	RV a ("{a }"); 
	RV b ("{b }"); RV ab = a; ab.add(b);
	RV c ("{c }"); RV ac = a; ac.add(c);
	RV r ("{r }");
	
	double at = 1.5;
	double bt = 0.5;
	double ct = -0.5;
	double rt = 0.30;
	// Full system
	vec thy =vec_2 ( at,bt ); //Simulated system - zero for constant term
	vec thz =vec_2 ( at,ct ); //Simulated system - zero for constant term
	vec Thy = concat ( thy, vec_1(rt) ); //Full parameter
	vec Thz = concat ( thz, vec_1(rt) ); //Full parameter

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

	ARX P1;
	P1.set_statistics(2, V0, -1 );
	ARX P2;
	 P2.set_statistics(2, V0, -1 );

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

	// Logging
	dirfilelog L ( "exp/merg_2a",3 );
	int Li_Data = L.add ( RV ( "{U Y Z }" ), "" );
//	int Li_LL   = L.add ( RV ( "{P1 P2 M1 M2 }" ), "LL" );
	int Li_P1m   = L.add ( concat ( ab,r ), "P1m" );
	int Li_P2m   = L.add ( concat ( ac,r ), "P2m" );
	int Li_Mm   = L.add ( concat ( ab,concat(r,c) ), "Mm" );
	int Li_Th   = L.add ( concat ( ab,concat(c,r) ), "T" );
	L.init();

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

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

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

	vec rgru(2);
	vec rgry(2);
	Yt(0) = 0.1;
	Ut(0) = 0.0;
	for ( t=1; t<ndat; t++ ) {
		// True system
		Ut ( t ) = pow ( sin ( ( t/40.0 ) *pi ),3 );
		rgru(0) = Ut(t); rgru(1) = Ut(t-1);
		Yt ( t ) = thy*rgru + rt * NorRNG();
		rgry(0) = Yt(t); rgry(1) = Yt(t-1);
		Zt ( t ) = thz*rgry + rt * NorRNG();

		// Bayes for all
		P1.bayes ( concat ( Yt ( t ),rgru ) );
		P2.bayes ( concat ( Zt ( t ),rgry ) );

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

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

		//Logger
		L.logit(Li_Data, vec_3(Ut(t), Yt(t), Zt(t)));
		L.logit(Li_P1m, P1._e()->mean());
		L.logit(Li_P2m, P2._e()->mean());
		L.logit(Li_Mm, M.mean());
		L.logit(Li_Th, concat(thy,vec_2(ct,rt)));
		L.step ( );

	}
	L.finalize( );
	L.itsave ( "merg_2a.it" );
	cout << endl;
}
