function [mu,classPi,z,si2,CLL,ILL,score] = gmm(X,K,init)
% input X is N*D data, K is number of clusters desired; numIter is the
% maximum number of iterations to run.  the output is: mu, which is K*D, the
% coordinates of the means; classPi, which is K*1 and represents the class prior
% probabilities; z, which is N*K, stores in positions z(n,k) the probability
% that the nth data point belongs to cluster k.  specifying the cluster
% associated with each data point; si2 is the estimated (shared) variance of
% the data; CLL is a plot of the complete log likelihoods and ILL is a plot
% of the incomplete log likelihoods.
  
[N D] = size(X);

if K >= N,
  error('gmm: you are trying to make too many clusters!');
end;

numIter = 20;  % maximum number of iterations to run

% initialize the log likelihoods
CLL = zeros(numIter,1);
ILL = zeros(numIter,1);

% initialize si2 dumbly
si2 = eye(D);

% initialize classPi uniformly
classPi = ones(K,1) / K;

% we initialize the means totally randomly
if nargin<3 || strcmp(init, 'random'),
  % initialize by choosing K (distinct!) random points: we do this by
  % randomly permuting the examples and then selecting the first K
  perm = randperm(N);
  perm = perm(1:K);
  
  % the centers are given by points
  mu = X(perm, :);
  
  % we leave the assignments unspecified -- they'll be computed in the
  % first iteration of K means
  z = zeros(N,1);

elseif strcmp(init, 'furthest'),
  % run furthest-first initialization
  mu = clusterInit(X,K);
  
  % again, don't bother initializing z
  z = zeros(N,1);
else
  error('unknown initialization: use "furthest" or "random"');
end;

for iter=1:20,
  % in the first step, we do assignments: each point is probabilistically
  % assigned to each center
  
  invSI2 = inv(si2 + 0.1 * eye(D));
  logdetSI2 = logdet(si2 + 0.1 * eye(D));
  
  for n=1:N,
      for k=1:K,
          % compute z(n,k) = log probability that the nth data point belongs
          % to cluster k... this is just like prediction from the
          % classification case
          
          z(n,k) = log(classPi(k)) -(1/2) * (log(2 * pi) + logdetSI2) + ...
                  -(1/2) * (X(n,:) - mu(k,:)) * invSI2 * (X(n,:) - mu(k,:))';
      end;

    % "magically" turn log probabilities into actual probabilities
    maxZ   = max(z(n,:));
    z(n,:) = exp(z(n,:) - maxZ - log(sum(exp(z(n,:) - maxZ))));
  end;
  
  % re-estimate classPi
  for k = 1:K
      classPi(k) = sum(z(:,k))/N;
  end;
  
  
  % re-estimate the means
  
  mu = zeros(K,D);
  for k=1:K
      for n = 1:N
          mu(k,:) = mu(k,:) + X(n,:) * z(n,k);
      end;
      mu(k,:) = mu(k,:)/sum(z(:,k));
  end;
 
  
  
  % re-estimate the variance
  si2 = zeros(D);
  for n=1:N,
    for k=1:K,
      si2 = si2 + z(n,k) * (X(n,:) - mu(k,:))' * (X(n,:) - mu(k,:))/N; 
    end;
    
  end;
  
  % compute *complete* and *incomplete* log likelihoods
  invSI2 = inv(si2 + 0.1 * eye(D));
  logdetSI2 = logdet(si2 + 0.1 * eye(D));
  cll = 0;
  ill = 0;
  for n=1:N,
    ilk = -Inf;
    for k=1:K,
      % compute pnk = joint log probability of data point x and cluster
      % k... note that this is exactly the same as z(n,k) from
      % above... feel free to just copy!
      pnk = log(classPi(k)) -(1/2) * (log(2 * pi) + logdetSI2) + ...
                  -(1/2) * (X(n,:) - mu(k,:)) * invSI2 * (X(n,:) - mu(k,:))';

      % update cll and ill
      cll = cll + z(n,k) * pnk;
      ilk = log(exp(ilk)+exp(pnk));
    end;
    ill = ill + ilk;
  end;
  
  % store values
  CLL(iter) = cll;
  ILL(iter) = ill;
end;

score = CLL(end);

function y = logdet(A)
  U = chol(A);
  y = 2*sum(log(diag(U)));
