%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Name: M_step.m
% Input: x - a nxd matrix (nx3 if using RGB)
%        Q - vector of values from the complete data log-likelihood function
%        h - a nxk matrix, the expectation of the hidden variable z given the data set and distribution params
%        S - cluster covariance matrices
%        k - the number of clusters
%        flag - flag to use improved EM to avoid singular covariance matrix
% Output: S - cluster covariance matrices
%         m - cluster means
%         pi - mixing coefficients
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [S, m, pi] = M_step(x, h, S, k, flag)

  % get size of data
  [num_data, dim] = size(x);
  eps = 1e-15;
  lambda = 1e-3; % value for improved version of EM

  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  % update mixing coefficients
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  pi = zeros(k, 1);
  N_i = zeros(k, 1);
  for i = 1:k
    N_i(i) = sum(h(:, i));
  end
  pi = N_i / num_data;
  
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  % update cluster means
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  m = zeros(k, dim);
  m = h' * x ./ N_i;
  % for i = 1:k
  %   m(i, :) = sum(h(:, i) .* x(i, :)) / N_i(i);
  % end

  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  % Calculate the covariance matrix estimate
  % further modifications will need to be made when doing 2(d)
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  S = zeros(dim, dim, k);
  for i = 1:k
    % for j = 1:num_data
    %   S(:, :, i) = S(:, :, i) + h(j, i) * (x(j, :) - m(i, :)) * (x(j, :) - m(i, :))';
    % end
    s = (x - m(i, :))' * ((x - m(i, :)) .* h(:, i)) / N_i(i);
    
    % % MAKE IT SYMMETRIC https://stackoverflow.com/a/38730499
    % S(:, :, i) = (s + s') / 2;
    % https://www.mathworks.com/matlabcentral/answers/366140-eig-gives-a-negative-eigenvalue-for-a-positive-semi-definite-matrix#answer_290270
    s = (s + s') / 2;
    [V, D] = eig(s);
    S(:, :, i) = V * max(D,eps) / V;
  end

end