#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 ), empty_vec );
		//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 mat &cond, 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_weighted ( data.get_col ( t ), empty_vec, 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, const vec &cond=empty_vec ) {

};

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


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

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

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

void MixEF::flatten ( const BMEF* M2 ) {
	const MixEF* Mix2 = dynamic_cast<const MixEF*> ( M2 );
	bdm_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 ) );
}
}