#include <vector>
#include "mixtures.h"

namespace bdm{


void MixEF::init ( BMEF* Com0, const mat &Data, int c ) {
	//prepare sizes
	Coms.set_size ( c );
	n=c;
	weights.set_parameters ( ones ( c ) ); //assume at least one observation in each comp.
	//est will be done at the end
	//
	int i;
	int ndat = Data.cols();
	//Estimate  Com0 from all data
	Coms ( 0 ) = Com0->_copy_();
//	Coms(0)->set_evalll(false);
	Coms ( 0 )->bayesB ( Data );
	// Flatten it to its original shape
	Coms ( 0 )->flatten ( Com0 );

	//Copy it to the rest
	for ( i=1;i<n;i++ ) {
		//copy Com0 and create new rvs for them
		Coms ( i ) =  Coms ( 0 )->_copy_ ( );
	}
	//Pick some data for each component and update it
	for ( i=0;i<n;i++ ) {
		//pick one datum
		int ind=floor(ndat*UniRNG.sample());
		Coms ( i )->bayes ( Data.get_col ( ind ),1.0 );
		//flatten back to oringinal
		Coms(i)->flatten(Com0);
	}

	//est already exists - must be deleted before build_est() can be used
	delete est;
	build_est();

}

void MixEF::bayesB ( const mat &data , const vec &wData ) {
	int ndat=data.cols();
	int t,i,niter;
	bool converged=false;

	multiBM weights0 ( weights );

	Array<BMEF*> Coms0 ( n );
	for ( i=0;i<n;i++ ) {Coms0 ( i ) = ( BMEF* ) Coms ( i )->_copy_();}

	niter=0;
	mat W=ones ( n,ndat ) / n;
	mat Wlast=ones ( n,ndat ) / n;
	vec w ( n );
	vec ll ( n );
	// tmp for weights
	vec wtmp = zeros ( n );
	int maxi;
	double maxll;
	//Estim
	while ( !converged ) {
		// Copy components back to their initial values
		// All necessary information is now in w and Coms0.
		Wlast = W;
		//
		//#pragma omp parallel for
		for ( t=0;t<ndat;t++ ) {
			//#pragma omp parallel for
			for ( i=0;i<n;i++ ) {
				ll ( i ) =Coms ( i )->logpred ( data.get_col ( t ) );
				wtmp =0.0; wtmp ( i ) =1.0;
				ll ( i ) += weights.logpred ( wtmp );
			}
			
			maxll = max(ll,maxi);
			switch (method) {
				case QB:
					w = exp ( ll-maxll );
					w/=sum(w);
					break;
				case EM:
					w = 0.0;
					w(maxi) = 1.0;
					break;
			}
			
			W.set_col ( t, w );
		}

		// copy initial statistics
		//#pragma omp parallel for
		for ( i=0;i<n;i++ ) {
			Coms ( i )-> set_statistics ( Coms0 ( i ) );
		}
		weights.set_statistics ( &weights0 );

		// Update statistics
		// !!!!    note  wData ==> this is extra weight of the data record
		// !!!!    For typical cases wData=1.
		for ( t=0;t<ndat;t++ ) {
			//#pragma omp parallel for
			for ( i=0;i<n;i++ ) {
				Coms ( i )-> bayes ( data.get_col ( t ),W ( i,t ) * wData ( t ) );
			}
			weights.bayes ( W.get_col ( t ) * wData ( t ) );
		}

		niter++;
		//TODO better convergence rule.
		converged = (niter>10);//( sumsum ( abs ( W-Wlast ) ) /n<0.1 );
	}

	//Clean Coms0
	for ( i=0;i<n;i++ ) {delete Coms0 ( i );}
}

void MixEF::bayes ( const vec &data ) {

};

void MixEF::bayes ( const mat &data ) {
	this->bayesB ( data, ones ( data.cols() ) );
};


double MixEF::logpred ( const vec &dt ) const {

	vec w=weights.posterior().mean();
	double exLL=0.0;
	for ( int i=0;i<n;i++ ) {
		exLL+=w ( i ) *exp ( Coms ( i )->logpred ( dt ) );
	}
	return log ( exLL );
}

mixtures* MixEF::epredictor (  ) const {
	Array<epdf*> pC ( n );
	for ( int i=0;i<n;i++ ) {pC ( i ) =Coms ( i )->epredictor ( );}
	mixtures* tmp;
	tmp = new mixtures (  );
	tmp->set_parameters ( weights.posterior().mean(), pC, false );
	tmp->ownComs();
	return tmp;
}

void MixEF::flatten(const BMEF* M2){
	const MixEF* Mix2=dynamic_cast<const MixEF*>(M2);
	it_assert_debug(Mix2->n==n,"Different no of coms");
	//Flatten each component
	for (int i=0; i<n;i++){
		Coms(i)->flatten(Mix2->Coms(i));
	}
	//Flatten weights = make them equal!!
	weights.set_statistics(&(Mix2->weights));
}
}
