# Tools for learning Gaussian mixture models.
# Adapted from the BRMLtoolkit by David Barber.

import numpy as np
import numpy.matlib as matlib
import numpy.linalg as LA
import scipy.misc

def condp(X):
	return X / np.sum(X, axis=0)

def condexp(logp):
	pmax = np.max(logp, axis=0)
	P = logp.shape[0]
	return condp(np.exp(logp - np.tile(pmax, (P, 1))))

def GMMlogp(X, H, P, m, S):
	D, N = X.shape  # dimension and number of points
	
	logp = np.zeros((N, H))
	for i in range(H):
		invSi = LA.inv(S[i,:,:])
		sign, logdetSi = LA.slogdet(2 * np.pi * S[i,:,:])
		
		for n in range(N):
			v = X[:,n] - m[:,i]
			logp[n,i] = -0.5 * (v @ invSi @ v) - 0.5 * logdetSi + np.log(P[i])
	
	return logp



# Log Likelihood of data X under a Gaussian Mixture Model
#
# X : each column of X is a datapoint.
# P : mixture coefficients
# m : means
# S : covariances
#
# Returns: A list containing the log likelihood for each data point in X
def GMMloglik(X, P, m, S):
	N = X.shape[1]
	H = m.shape[1]
	
	logp = GMMlogp(X, H, P, m, S)
	logl = [scipy.misc.logsumexp(a=logp[n,:], b=np.ones(H)) for n in range(N)]
	
	return logl

# Fit a mixture of Gaussian to the data X using EM
#
# X : each column of X is a datapoint.
# H : number of components of the mixture.
# n_iter : number of EM iterations
#
# Returns: (P, m, S, loglik, phgn)
# P : learned mixture coefficients
# m : learned means
# S : learned covariances
# loglik : log likelihood of the learned model
# phgn : mixture assignment probabilities
def GMMem(X, H, n_iter):
	D, N = X.shape  # dimension and number of points
	
	# initialise the centres to random datapoints
	r = np.random.permutation(N)
	m = X[:, r[:H]]
	
	# initialise the variances to be large
	s2 = np.mean(np.diag(np.cov(X)))
	S = matlib.tile(s2 * np.eye(D), [H, 1, 1])
	
	# intialise the component probilities to be uniform
	P = np.ones(H) / H
	
	for emloop in range(n_iter):
		# E-step:
		logpold = GMMlogp(X, H, P, m, S)
		
		phgn = condexp(logpold.T)  # responsibilities
		pngh = condp(phgn.T)       # membership
		
		# M-step:
		for i in range(H):   # now get the new parameters for each component
			tmp = (X - np.tile(m[:,i:i+1], N)) * np.tile(np.sqrt(pngh[:,i]), (D,1))
			Scand = np.dot(tmp, tmp.T)
			
			if LA.det(Scand) > 0.0001:   # don't accept too low determinant
				S[i,:,:] = Scand
		
		m = np.dot(X, pngh)
		P = np.sum(phgn, axis=1) / N
	
	logl = np.sum(scipy.misc.logsumexp(a=logpold[n,:], b=np.ones(H)) for n in range(N))
	
	return P, m, S, logl, phgn