function [params] = gaussian_mixture_mstep(data,memberships);
% function [params] = gaussian_mixture_mstep(data,memberships); 
%
%  Estimates new parameters (M-step) for a Gaussian mixture model with d
%  variables, given membership probabilities from the E-step
%
%  INPUTS
%  data: N x d real-valued data matrix
%  memberships = n x K matrix of component memberships
%
%  OUTPUTS  
%  params:  2d array of structures of the form: 
%           params(k).weight = weight of component k
%           params(k).mean = d-dimensional mean vector for kth component 
%           params(k).covariance = d x d covariance vector for kth
%           component 
%  
%                                        Padhraic Smyth, CS 274A, Winter 2009 

% Ideally need to add some error-checking here....

[n d] = size(data);
[n K] = size(memberships);
c = cov(data);

alpha = 0.0001;  % alpha * sigma_j^2 is the lowest allowable value for the
              % the variance of any component


% calculate maximum likelihood parameter estimates given
% probabilistic memberships (for each component and each variable)
for k=1:K
   % get the vector of weights (one weight per row) for component k
   w = memberships(:,k);
   sumw = sum(w);
   wmatrix = repmat(w,1,d);
   
   % estimate the mean of component k: 
   params(k).mean = sum(data.*wmatrix)/sumw;
   
   % estimate the covariance matrices of component k:
   mean_matrix = repmat(params(k).mean,n,1); 
   delta = data-mean_matrix;
   
   % d^2 for-loop below - it may be possible to "matricize" this
   for i=1:d
       for j=i:d
           params(k).covariance(i,j) = sum(delta(:,i).*delta(:,j).*w)/sumw;
           params(k).covariance(j,i) = params(k).covariance(i,j);
       end
   end
   %delta_matrix = (data-mean_matrix).*(data-mean_matrix);
   %params(k).variances = sum(delta_matrix.*wmatrix)/sumw; 
   for j=1:d
        if params(k).covariance(j,j)<(alpha*c(j,j))
             params(k).covariance(j,j) = alpha*c(j,j); 
        end 
   end
   
   % estimate the relative weight of each component
   params(k).weight = sum(w)./n;
end