#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;

/*! \file
Experiment for distributed identification using log-normal merging

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:
\dot
digraph compart{

  edge [fontname="FreeSans",fontsize=10,labelfontname="FreeSans",labelfontsize=10];
  node [fontname="FreeSans",fontsize=10,shape=record];

rankdir=LR;

 U [label="u",height=0.2,width=0.4,color="white", fillcolor="white", style="filled" fontcolor="black"];
 AR1 [label="a,b,r",height=0.2,width=0.4,color="black", fillcolor="white", style="filled" fontcolor="black"]
 U -> AR1 [color="midnightblue",style="solid"]; 
 AR2 [label="a,c,r",height=0.2,width=0.4,color="black", fillcolor="white", style="filled" fontcolor="black"]
 AR1 -> AR2 [color="midnightblue",style="solid",label="y"]; 
 Z [label="z",height=0.2,width=0.4,color="white", fillcolor="white", style="filled" fontcolor="black"];
 AR2 -> Z [color="midnightblue",style="solid"]; 

 
}
\enddot 
*/

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 = 0.9;
	double bt = 5;
	double ct = -0.5;
	double rt = 0.50;
	// 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.01*eye ( 3 ); V0 ( 0,0 ) = 1; //

	ARX P1; P1.set_rv(concat(ab,r));
	P1.set_statistics(1, V0, -1 );
	P1.set_parameters(0.9);
	ARX P2; P2.set_rv(concat(ac,r));
	P2.set_statistics(1, V0, -1 );
	P2.set_parameters(0.9);

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

	// Logging
	dirfilelog L ( "exp/merg_2a",100 );
	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_rv(ab); 
	g0.set_parameters ( "1 1 ",mat ( "1 0; 0 1" ) );
	egamma g1;g1.set_rv(r); 
	g1.set_parameters ( "2  ", "2" );
	enorm<ldmat> g2; g2.set_rv(c); 
	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 );

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

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

		if (t>50) {bt+=0.1;
			thy(1)=bt;}
		if (t>50) {ct-=0.01;
			thz(1)=ct;}
		
		// Merge estimates
		Array<pdf*> A ( 2 ); A ( 0 ) =&P1;A ( 1 ) =&P2;
		merger_mix M ( A );
		M.set_parameters ( 20 ,0.99); 
		M.set_method(LOGNORMAL, 1.2);
		M._Mix().set_method(QB);
		//M2.set_parameters ( 100.0, 1000,3 ); //M2._Mix().set_method(QB);
/*		char fnm[100];
		sprintf(fnm,"m2a_dbg%d.it",t);
		M.debug_file(fnm);*/
		M.set_support ( *proposal,100 );
		M.merge();
		//proposal = M.proposal();
		//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*rt)));
		L.step ( );

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