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

namespace bdm {


void MixEF::init ( BMEF* Com0, const mat &Data, const int c ) {
	//prepare sizes
	Coms.set_size ( c );
	weights.set_parameters ( ones ( c ) ); //assume at least one observation in each comp.
	multiBM weights0(weights);
	//est will be done at the end
	//
	int i;
	int ndat = Data.cols();
	//Estimate  Com0 from all data
	Coms ( 0 ) = (BMEF*) Com0->_copy();
//	Coms(0)->set_evalll(false);
	Coms ( 0 )->bayes_batch ( Data );

	Coms ( 0 )->flatten ( Com0 ); 

	//Copy it to the rest
	for ( i = 1; i < Coms.length(); i++ ) {
		//copy Com0 and create new rvs for them
		Coms ( i ) =  (BMEF*) Coms ( 0 )->_copy ( );
	}
	//Pick some data for each component and update it
	for ( i = 0; i < Coms.length(); i++ ) {
		//pick one datum
		if (ndat==Coms.length()){ //take the ith vector
			Coms ( i )->bayes ( Data.get_col ( i ), empty_vec );
		} else { // pick at random
			int ind = (int) floor ( ndat * UniRNG.sample() );
			Coms ( i )->bayes_weighted ( Data.get_col ( ind ), empty_vec, ndat );
			Coms (i)->flatten(Com0,ndat/Coms.length());
		}
		//sharpen to the sharp component
		//Coms ( i )->flatten ( SharpCom.get(), 1.0/Coms.length() );
	}
	MixEF_options old_opt =options;
	MixEF_options ini_opt=options;
	ini_opt.method = EM;
	ini_opt.max_niter= 1;
	bayes_batch(Data, empty_vec);
	
	for ( i = 0; i < Coms.length(); i++ ) {
		Coms (i)->flatten(Com0,ndat/Coms.length());
	}
	
	options = old_opt;
}

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

	multiBM weights0 ( weights );

	int n = Coms.length();
	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;
	
	double levid=0.0;
	//Estim
	while ( !converged ) {
		levid=0.0;
		// 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 ) , empty_vec);
				wtmp = 0.0;
				wtmp ( i ) = 1.0;
				ll ( i ) += weights.logpred ( wtmp );
			}

			maxll = max ( ll, maxi );
			switch ( options.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.
		vec logevid(n);
		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 ) );
				logevid(i) = Coms(i)->_ll();
			}
			weights.bayes ( W.get_col ( t ) * wData ( t ) );
		}
		levid += weights._ll()+log(weights.posterior().mean() * exp(logevid)); // inner product w*exp(evid)
		
		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 );
	}
	return levid;
}

void MixEF::bayes ( const vec &data, const vec &cond = empty_vec ) {

};

double MixEF::logpred ( const vec &yt, const vec &cond =empty_vec) const {

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

emix* MixEF::epredictor ( const vec &vec) const {
	Array<shared_ptr<epdf> > pC ( Coms.length() );
	for ( int i = 0; i < Coms.length(); i++ ) {
		pC ( i ) = Coms ( i )->epredictor ( );
		pC (i) -> set_rv(_yrv());
	}
	emix* tmp;
	tmp = new emix( );
	tmp->_w() = weights.posterior().mean();
	tmp->_Coms() = pC;
	tmp->validate();
	return tmp;
}

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