function [rho,mu,Psi,r,lik] = mog(X,K,rho,mu,Psi,r);
% function [rho,mu,Psi,r,lik] = mog(X,K,rho,mu,Psi,r);
% 
% 1 iteration of EM for mixture of Gaussians.
% inputs rho, mu, Psi, r optional; initialized randomly if not given.
%
% X(i,:)        data item i
% K             number of Gaussian clusters
% m             number of data items
% n             dimensionality of data items
% rho(1,k)      mixing proportion of cluster k
% mu(k,:)       mean of cluster k
% Psi(:,:,k)    covariance of cluster k
% r(i,k)        responsibility of cluster k for data item i
% lik           likelihood at current parameter settings

epsilon = 1e-3; % regularizer to ensure EM don't blow up

[m n] = size(X);
S = cov(X);

if nargin == 2
  % initialize parameters
  rho = 1+rand(1,K);
  rho = rho/sum(rho);
  c = randperm(m);
  mu = X(c(1:K),:);
  Psi = S(:,:,ones(1,K))/1000;
  r = zeros(m,K);
end

% E step
lik = 0;
const = zeros(1,K);
p = zeros(1,K);
sq = zeros(n,n,K);
for c = 1:K
  sq(:,:,c) = chol(Psi(:,:,c));
  const(c) = log(rho(c))-n/2*log(2*pi)-sum(log(diag(sq(:,:,c))));
end
for i = 1:m
  for c = 1:K
    delta = (X(i,:)-mu(c,:))/sq(:,:,c);
    p(c) = const(c) - .5*(delta*delta');
  end
  r(i,:) = exp(p-max(p));
  r(i,:) = r(i,:) / sum(r(i,:));
  lik = lik + r(i,:)*p' - r(i,:)*log(r(i,:)+1e-100)';
end

% M step
for c = 1:K
  rho(c) = (sum(r(:,c),1) + epsilon)/(m+epsilon);
  mu(c,:) = r(:,c)'*X / sum(r(:,c));
  delta = X - mu(c*ones(1,m),:);
  Psi(:,:,c) = ((r(:,c*ones(1,n)).*delta)'*delta + S*epsilon) ...
                / (sum(r(:,c))+epsilon);
end
