#ifdef NICE_USELIB_OPENMP
#include <omp.h>
#endif

#include <stdio.h>

#include "GMM.h"
#include <core/vector/Algorithms.h>

#include "vislearning/math/cluster/KMeans.h"

#define DEBUGGMM

using namespace OBJREC;
using namespace NICE;
using namespace std;

GMM::GMM()
{
	gaussians = 3;
	maxiter = 200;
	featsperclass = 0;
	tau = 10.0;
	srand ( time (NULL) );
	comp = false;
}

GMM::GMM(int _gaussians):gaussians(_gaussians)
{
	maxiter = 200;
	featsperclass = 0;
	tau = 0.0;
	srand ( time (NULL) );
	pyramid = false;
	comp = false;
}

GMM::GMM(const Config *conf, int _gaussians):gaussians(_gaussians)
{
	this->conf = conf;
	
	if(gaussians < 2)
		gaussians = conf->gI("GMM", "gaussians", 2 );
		
	maxiter = conf->gI("GMM", "maxiter", 200 );
		
	featsperclass = conf->gI("GMM", "featsperclass", 0 );
	
	tau = conf->gD("GMM", "tau", 100.0);
	
	pyramid = conf->gB("GMM", "pyramid", false);
	
	comp = false;
	
	srand ( time (NULL) );
}

void GMM::computeMixture(Examples examples)
{
	int fsize = (int)examples.size();
	assert(fsize >= gaussians);
	
	dim = examples[0].second.vec->size();
	
	int samples = fsize;
	if(featsperclass > 0)
	{
		samples = featsperclass*gaussians;
		samples = std::min(samples, fsize);
	}
		
	// Copy data in Matrix
	VVector dataset;

	cout << "reduced training data for GMM from " << fsize << " features to " << samples << " features.";
	for(int i = 0; i < samples; i++)
	{
		int k = rand() % fsize;
		NICE::Vector *vec = examples[k].second.vec;
		dataset.push_back(*vec);
	}
	
	computeMixture(dataset);
}

void GMM::computeMixture(const VVector &DataSet)
{
	// Learn the GMM model
	assert(DataSet.size() >= (uint)gaussians);
	//initEMkMeans(DataSet); // initialize the model
	srand ( time (NULL) );
	
	bool poweroftwo = false;
	int power = 1;
	while(power <= gaussians)
	{
		if(gaussians == power)
			poweroftwo = true;
		power *= 2;
	}
	
	if(poweroftwo && pyramid)
	{
		initEM(DataSet); // initialize the model
		int g = 1;
		
		while(g <= gaussians)
		{
			cout << "g = " << g << endl;
			doEM(DataSet, g);
			
			if(2*g <= gaussians)
			{
				for(int i = g; i < g*2; i++)
				{
					mu[i] = mu[i-g];
					sparse_sigma[i] = sparse_sigma[i-g];
					sparse_inv_sigma[i] = sparse_inv_sigma[i-g];
					log_det_sigma[i] = log_det_sigma[i-g];
					priors[i] = priors[i-g];
					
					double interval = 1.0;
					for(int k = 0; k < (int)mu[i].size(); k++)
					{
						interval = mu[i][k];
						interval = std::max(interval, 0.1);
						double r = (interval*((double)rand()/(double)RAND_MAX))-interval/2.0;
						mu[i][k] += r;
					}
				}
			}
			g *= 2;
		}
	}
	else
	{
		initEMkMeans(DataSet); // initialize the model
		doEM(DataSet, gaussians);
	}
	// performs EM
}

inline double diagDeterminant(const NICE::Vector &sparse_mat)
{
	double det = 1.0;
	for(int i = 0; i < (int)sparse_mat.size(); i++)
	{
		det *= sparse_mat[i];
	}
	
	return det;
}

inline double logdiagDeterminant(const NICE::Vector &sparse_mat)
{
	double det = 0.0;
	for(int i = 0; i < (int)sparse_mat.size(); i++)
	{
		det += log(sparse_mat[i]);
	}

	return det;
}

inline NICE::Vector diagInverse(const NICE::Vector &sparse_mat)
{
	NICE::Vector inv(sparse_mat.size());
	for(int i = 0; i < (int)sparse_mat.size(); i++)
	{
		inv[i] = 1.0/sparse_mat[i];
	}
	return inv;
}

void GMM::initEMkMeans(const VVector &DataSet)
{
	/*init GMM with k-Means*/
	KMeans k( gaussians );
	VVector means;
	vector<double> weights;
	vector<int> assignment;
	k.cluster ( DataSet, means, weights, assignment );
	
	int nData = DataSet.size();
	this->dim = DataSet[0].size();
	cdimval = dim*log(2*3.14159);
	vector<int> pop(gaussians, 0);
	priors.resize(gaussians);
	mu = VVector(gaussians, dim);
	log_det_sigma.clear();
	vector<int> allk;
	
	NICE::Vector globmean(dim);
	globmean.set(0.0);
	
	for(int n = 0; n < (int)DataSet.size(); n++) /* getting the max value for time */
	{
		globmean = globmean + DataSet[n];
	}

	globmean *= (1.0/nData);
	
	NICE::Vector sparse_globsigma;

	sparse_globsigma.resize(dim);
	sparse_globsigma.set(0.0);
	
	for(int i = 0; i < (int)DataSet.size(); i++) // Covariances updates
	{
		// nur diagonal Elemente berechnen
		for(int d = 0; d < dim; d++)
		{
			double diff = (DataSet[i][d] - globmean[d]);
			sparse_globsigma[d] += diff*diff;
		}
	}
	
	sparse_globsigma *= (1.0/DataSet.size());
	
	for(int i = 0; i < gaussians; i++)
	{
		NICE::Vector _inv_sigma = diagInverse(sparse_globsigma);

		sparse_sigma.push_back(sparse_globsigma);
		sparse_inv_sigma.push_back(_inv_sigma);
		log_det_sigma.push_back(logdiagDeterminant(sparse_globsigma));
		mu[i] = means[i];
		//priors[i]=1.0/(double)gaussians; // set equi-probables states 
		priors[i]=weights[i]; // set kMeans weights 
	}
}

void GMM::initEM(const VVector &DataSet)
{

	/* init the GaussianMixture by using randomized meanvectors */
	int nData = DataSet.size();
	this->dim = DataSet[0].size();
	cdimval = dim*log(2*3.14159);
	vector<int> pop(gaussians, 0);
	priors.resize(gaussians);
	mu = VVector(gaussians, dim);
	log_det_sigma.clear();
	vector<int> allk;
	
	NICE::Vector globmean(dim);
	globmean.set(0.0);
	
	for(int n = 0; n < (int)DataSet.size(); n++) /* getting the max value for time */
	{
		globmean = globmean + DataSet[n];
	}

	globmean *= (1.0/nData);
	
	NICE::Vector sparse_globsigma;

	sparse_globsigma.resize(dim);
	sparse_globsigma.set(0.0);
	
	for(int i = 0; i < (int)DataSet.size(); i++) // Covariances updates
	{
		// nur diagonal Elemente berechnen
		for(int d = 0; d < dim; d++)
		{
			double diff = (DataSet[i][d] - globmean[d]);
			sparse_globsigma[d] += diff*diff;
		}
	}
	
	sparse_globsigma *= (1.0/DataSet.size());
	
	for(int i = 0; i < gaussians; i++)
	{
		priors[i]=1.0/(double)gaussians; // set equi-probables states 
		NICE::Vector _inv_sigma = diagInverse(sparse_globsigma);

		sparse_sigma.push_back(sparse_globsigma);
		sparse_inv_sigma.push_back(_inv_sigma);
		log_det_sigma.push_back(logdiagDeterminant(sparse_globsigma));
		bool newv = false;
		int k;
		while(!newv)
		{
			newv = true;
			k = rand() % nData;
		
			for(int nk = 0; nk < (int)allk.size(); nk++)
				if(allk[nk] == k)
				{
					newv = false;
				}
			if(newv)
				allk.push_back(k);
		}
		mu[i] = DataSet[k];
	}
}

inline void sumRow(NICE::Matrix mat, NICE::Vector &res)
{
	int row = mat.rows();
	int column = mat.cols();
	res =NICE::Vector(column);
	res.set(1e-5f);
//#pragma omp parallel for
	for(int i = 0; i < column; i++){
		for(int j = 0; j < row; j++){
			res[i] += mat(j, i);
		}
	}
}

double GMM::logpdfState(const NICE::Vector &Vin,int state)
{
	/* get the probability density for a given state and a given vector */
	double p;
	NICE::Vector dif;
		
	dif = Vin;
	dif -= mu[state];
		
	p = 0.0;
	for(int i = 0; i < (int)dif.size(); i++)
	{
		p += dif[i] * dif[i] * sparse_inv_sigma[state][i];
	}
	
	p = -0.5*(p + cdimval + log_det_sigma[state]);
	return p;
}

int GMM::doEM(const VVector &DataSet, int nbgaussians)
{
  /* perform Expectation/Maximization on the given Dataset :
	Matrix DataSet(nSamples,Dimensions).
	The GaussianMixture Object must be initialised before 
	(see initEM or initEMkMeans methods ) */
	
#ifdef DEBUG
	cerr << "GMM::start EM" << endl;
#endif	
	
	int nData = DataSet.size();
	int iter = 0;
	double log_lik;
	double log_lik_threshold = 1e-6f;
	double log_lik_old = -1e10f;

	NICE::Matrix unity(dim,dim,0.0);
	for(int k = 0; k < dim; k++) 
		unity(k, k)=1.0;
	
	//EM loop
	while(true)
	{
#ifdef DEBUGGMM
		cerr << "GMM::EM: iteration: " << iter << endl;
#endif
		
		vector<double> sum_p;
		sum_p.resize(nData);

		for(int i = 0; i < nData; i++)
		{
			sum_p[i] = 0.0;
		}
	
		NICE::Matrix pxi(nData,gaussians);
		pxi.set(0.0);
		NICE::Matrix pix(nData,gaussians);
		pix.set(0.0);
		NICE::Vector E;
		
		iter++;
		if (iter>maxiter){
			cerr << "EM stops here. Max number of iterations (" << maxiter << ") has been reached." << endl;
			return iter;
		}

		double sum_log = 0.0; 
		
    	// computing expectation
		double maxp = -numeric_limits<double>::max();
		vector<double> maxptmp(nData,-numeric_limits<double>::max());
#pragma omp parallel for
		for(int i = 0; i < nData; i++)
		{
			for(int j = 0; j < nbgaussians; j++)
			{
				double p = logpdfState(DataSet[i],j); // log(P(x|i))
				maxptmp[i] = std::max(maxptmp[i], p);
				pxi(i, j) = p;
			}
		}
		
		for(int i = 0; i < nData; i++)
		{
			maxp = std::max(maxp, maxptmp[i]);
		}

#pragma omp parallel for
		for(int i = 0; i < nData; i++)
		{
			sum_p[i] = 0.0;
			for(int j = 0; j < nbgaussians; j++)
			{
				double p = pxi(i, j) - maxp; // log(P(x|i))
				p = exp(p); // P(x|i)
				
				if(p < 1e-40)
					p = 1e-40;

				pxi(i, j) = p;
				sum_p[i] += p*priors[j];
			} 
		}
			
		for(int i = 0; i < nData; i++)
		{
			sum_log += log(sum_p[i]);
		}
		
#pragma omp parallel for
		for(int j = 0; j < nbgaussians; j++)
		{
			for(int i = 0; i < nData; i++)
			{
				pix(i, j) = (pxi(i, j)*priors[j])/sum_p[i]; // then P(i|x)			
			}   
		}

		// here we compute the log likehood 
		log_lik = sum_log/nData; 
#ifdef DEBUGGMM	
		cout << "diff: " << fabs((log_lik/log_lik_old)-1) << " thresh: " << log_lik_threshold << " sum: " << sum_log <<  endl;
		//cout << "logold: " << log_lik_old << " lognew: " << log_lik << endl;
#endif
		
		if(fabs((log_lik/log_lik_old)-1) < log_lik_threshold )
		{
			//if log likehood hasn't move enough, the algorithm has converged, exiting the loop 
			return iter;
		}

		log_lik_old = log_lik;

		// Update Step
		sumRow(pix,E);
		
#pragma omp parallel for
		for(int j = 0; j < nbgaussians; j++) 
		{
			priors[j] = (E[j]+tau)/(nData+tau*nbgaussians); // new priors 

			NICE::Vector tmu(dim);
			tmu.set(0.0);

			NICE::Vector sparse_tmsigma(dim);
			sparse_tmsigma.set(0.0);

			for(int i = 0; i < nData; i++) // Means update loop
			{
				tmu = tmu + (DataSet[i]*pix(i,j));
			}
									
			NICE::Vector tmu2 = mu[j];
			mu[j] = tmu + tau*tmu2;
			mu[j] = mu[j]*(1.0/(E[j]+tau));
			
			for(int i = 0; i < nData; i++) // Covariances updates
			{	
				// nur diagonal Elemente berechnen
				for(int d = 0; d < dim; d++)
				{
					sparse_tmsigma[d] += DataSet[i][d]*DataSet[i][d]*pix(i, j);
				}
			}
			
			NICE::Vector sparse_tmsigma2(dim);
			for(int d = 0; d < dim; d++)
			{
				sparse_tmsigma[d] += 1e-5f;
				sparse_tmsigma2[d] = sparse_sigma[j][d]+tmu2[d]*tmu2[d];
				sparse_sigma[j][d] = (sparse_tmsigma[d]+tau*sparse_tmsigma2[d])/(E[j]+tau)-(mu[j][d]*mu[j][d]);
			}				
			sparse_inv_sigma[j] = diagInverse(sparse_sigma[j]);
				
			log_det_sigma[j] = logdiagDeterminant(sparse_sigma[j]);
		}
		if(comp)
		{
			double dist = compare();
			cout << "dist for iteration " << iter << endl;
			cout << "d: " << dist << endl;
		}	
	}
#ifdef DEBUG
	cerr << "GMM::finished EM after " << iter << " iterations" << endl;
#endif
	return iter;
}

int GMM::getBestClass(const NICE::Vector &v, double *bprob)
{
	int bestclass = 0;
	double maxprob = logpdfState(v,0)+log(priors[0]);  // log(P(x|i))
	
	for(int i = 1; i < gaussians; i++)
	{
		double prob = logpdfState(v,i)+log(priors[i]);  // log(P(x|i))
		
		if(prob > maxprob)
		{
			maxprob = prob;
			bestclass = i;
		}
	}
	
	if(bprob != NULL)
		*bprob = maxprob;
	
	return bestclass;
}
		
void GMM::getProbs(const NICE::Vector &vin, SparseVector &outprobs)
{
	outprobs.clear();
	outprobs.setDim(gaussians);
	Vector probs;
	getProbs(vin, probs);
	for(int i = 0; i < gaussians; i++)
	{
		if(probs[i] > 1e-7f )
			outprobs[i] = probs[i];
	}
}
		
void GMM::getProbs(const NICE::Vector &vin, Vector &outprobs)
{
	Vector probs;
	probs.resize(gaussians);
	probs.set(0.0);
		
	double probsum = 0.0;
	double maxp = -numeric_limits<double>::max();
	for(int i = 0; i < gaussians; i++)
	{	
		probs[i] = logpdfState(vin, i) + log(priors[i]);  // log(P(x|i))
		maxp = std::max(maxp, probs[i]);		
	}
		
	for(int i = 0; i < gaussians; i++)
	{
		probs[i] -= maxp;
		probs[i] = exp(probs[i]);
		probsum += probs[i];
	}

	// normalize probs
#pragma omp parallel for
	for(int i = 0; i < gaussians; i++)
	{
		probs[i] /= probsum;
	}
	outprobs = probs;
}

void GMM::getFisher(const NICE::Vector &vin, SparseVector &outprobs)
{
	int size = gaussians*2*dim;
	outprobs.clear();
	outprobs.setDim(size);
	
	int counter = 0;
	
	NICE::Vector classprobs;
	classprobs.resize(gaussians);
	classprobs.set(0.0);

	double maxp = -numeric_limits<double>::max();
	for(int i = 0; i < gaussians; i++)
	{
		classprobs[i] = logpdfState(vin, i) + log(priors[i]);  // log(P(x|i))

		maxp = std::max(maxp, classprobs[i]);
	}

	for(int i = 0; i < gaussians; i++)
	{
		double p = classprobs[i] - maxp;
		
		p = exp(p);
		
		for(int d = 0; d < dim; d++)
		{
			double diff = vin[d]-mu[i][d];
			
			double sigma2 = sparse_sigma[i][d]*sparse_sigma[i][d];
			double sigma3 = sigma2*sparse_sigma[i][d];
			double normmu = sqrt(priors[i]/sigma2);
			double normsig = sqrt((2.0*priors[i])/sigma2);
			
			double gradmu = (p * (diff/sigma2))/normmu;
			double gradsig = (p * ((diff*diff)/sigma3 - 1.0/sparse_sigma[i][d])) /normsig;
			if(fabs(gradmu) > 1e-7f)
				outprobs[counter] = gradmu;
			counter++;
			
			if(fabs(gradsig) > 1e-7f)
				outprobs[counter] = gradsig;
			counter++;
		}
	}
}

void GMM::cluster ( const VVector & features, VVector & prototypes, vector<double> & weights, vector<int> & assignment )
{
	computeMixture(features);

	prototypes.clear();
	weights.clear();
	assignment.clear();

	for(int i = 0; i < (int)features.size(); i++)
	{
		int c = getBestClass(features[i]);
		assignment.push_back(c);
		
		weights.push_back(priors[c]);
	}
	for(int c = 0; c < gaussians; c++)
		prototypes.push_back(mu[c]);
	
	cout << "tau: " << tau << endl;
}

void GMM::saveData(const std::string filename)
{
	ofstream fout( filename.c_str() );

	fout << gaussians << " " << dim << endl;
	
	mu >> fout;
	fout << endl;
			
	for(int n = 0; n < gaussians; n++)
	{
		fout << sparse_sigma[n] << endl;
	}
	
	for(int n = 0; n < gaussians; n++)
	{
		fout << priors[n] << " ";
	}

	fout.close();
}

bool GMM::loadData(const std::string filename)
{
	cerr << "read GMM Data" << endl;
	ifstream fin( filename.c_str() );

	if(fin.fail())
	{
		cerr << "GMM: Datei " << filename << " nicht gefunden!" << endl;
		return false;
	}

	fin >> gaussians;
	fin >> dim;
	cdimval = log(pow(2*3.14159,dim));
	mu.clear();

	for(int i = 0; i < gaussians; i++)
	{
		NICE::Vector tmp;
		fin >> tmp;
		mu.push_back(tmp);
	}
	
	log_det_sigma.clear();

	for(int n = 0; n < gaussians; n++)
	{
		NICE::Matrix _sigma;
		NICE::Vector _sparse_sigma;
		
		_sparse_sigma =NICE::Vector(dim);
		fin >> _sparse_sigma;
		sparse_sigma.push_back(_sparse_sigma);
				
		sparse_inv_sigma.push_back(diagInverse(sparse_sigma[n]));
		log_det_sigma.push_back(logdiagDeterminant(sparse_sigma[n]));
	}

	for(int n = 0; n < gaussians; n++)
	{
		double tmpd;
		fin >> tmpd;
		priors.push_back(tmpd);
	}
	fin.close();

	cerr << "reading GMM Data finished" << endl;
	return true;
}

void GMM::getParams(VVector &mean, VVector &sSigma, vector<double> &p)
{
	mean = mu;
	sSigma.resize(gaussians);
	p.clear();
	for(int i = 0; i < gaussians; i++)
	{
		sSigma[i] = sparse_sigma[i];
		p.push_back(priors[i]);
	}
}

void GMM::setCompareGM(VVector mean, VVector sSigma, vector<double> p)
{
	mu2 = mean;
	sparse_sigma2 = sSigma;
	priors2 = vector<double>(p);
}
		
double GMM::kPPK(NICE::Vector sigma1, NICE::Vector sigma2, NICE::Vector mu1, NICE::Vector mu2, double p)
{
	double d = mu1.size();
	
	double det1 = 1.0;
	double det2 = 1.0;
	double det3 = 1.0;
	double eval = 0.0;
	for(int i = 0; i < d; i++)
	{
		det1 *= sigma1[i];
		det2 *= sigma2[i];
		double sigma = 1.0/((1.0/sigma1[i]+1.0/sigma2[i])*p);
		det3 *= sigma;
		double mu = p*(mu1[i]*sigma1[i]+mu2[i]*sigma2[i]);
		eval += 0.5*mu*sigma*mu - (p*mu1[i]*mu1[i])/(2.0*sigma1[i]) - (p*mu2[i]*mu2[i]) / (2.0*sigma2[i]);
	}
	
	return (pow(2.0*M_PI, ((1.0-2.0*p)*d)/2.0) * pow(det1, -p/2.0) * pow(det2, -p/2.0) * pow(det3, 0.5) * exp(eval));
}

double GMM::compare()
{
	double distkij = 0.0;
	double distkjj = 0.0;
	double distkii = 0.0;
	for(int i = 0; i < gaussians; i++)
	{
		for(int j = 0; j < gaussians; j++)
		{
			double kij = kPPK(sparse_sigma[i],sparse_sigma2[j], mu[i], mu2[j], 0.5);
			double kii = kPPK(sparse_sigma[i],sparse_sigma[j], mu[i], mu[j], 0.5);
			double kjj = kPPK(sparse_sigma2[i],sparse_sigma2[j], mu2[i], mu2[j], 0.5);
			kij = priors[i]*priors2[j]*kij;
			kii = priors[i]*priors[j]*kii;
			kjj = priors2[i]*priors2[j]*kjj;
			distkij += kij;
			distkii += kii;
			distkjj += kjj;
		}
	}
	return (distkij / (sqrt(distkii)*sqrt(distkjj)));
}

void GMM::comparing(bool c)
{
	comp = c;
}