function [R_train, W, xbar, Psi, lls] = SparseFA_EM(S_train, D, max_iters)
%[R_train, W, xbar, Psi, lls] = SparseFA_EM(S_train, D, max_iters)
%  Input: 
%     S_train   - #movies X #users sparse matrix of movie ratings
%     D         - number of factors
%     max_iters - number of iterations to run EM
%  Output:
%     R_train   - low dimensional representation of each user (#users X D)
%     W         - factor loading matrix (#movies X D)
%     xbar      - mean ratings for each movie (#movies X 1)
%     Psi       - covariance matrix for residuals (vector of D diagonal entries)
%     lls       - marginal log-likelihood bound at each iteration

  [num_movies, num_users] = size(S_train);
  conv_thresh = 0;
  
  averageR = full(sum(S_train,2));
  for ind = 1:num_movies
    averageR(ind) = averageR(ind) / nnz(S_train(ind,:));
  end

  % initialize
  W = sqrt(10)*randn( [ num_movies, D ] );   % init transform
  Psi = ones(num_movies, 1);                 % covariance matrix
  xbar = averageR;                           % mean of observed ratings
  %xbar = zeros( [num_movies, 1] );           % mean observation

  % 
  % EXPECTATION MAXIMIZATION ALGORITHM
  % 
  done = false;  iters = 0;
  lls = zeros(max_iters);
  loglike = -1/eps;
  while ~done 

    iters = iters + 1;
    loglike_old = loglike;

    % init iteration
    Psi_new = zeros(num_movies, 1);
    Ez = zeros( D+1, num_users );
    Ezz = zeros( D+1, D+1, num_users );
    Q = zeros( D, D, num_users );
    Wtop = zeros( D+1, num_movies );  
    Wbottom = zeros( D+1, D+1, num_movies );

    %
    % E-STEP
    % 
    for n=1:num_users
      x_n = S_train( :, n );    
      I = find( x_n );
      if isempty(I), continue; end; 

      % construct observed quantities
      W0 = W( I, : );
      invPsi0 = diag(1./Psi(I));
      x0 = x_n( I );
      xbar0 = xbar( I );

      % compute sufficient statistics     
      Q(:,:,n) = inv( eye(D) + W0' * invPsi0 * W0 );
      m = Q(:,:,n) * (W0' * invPsi0 * ( x0 - xbar0 ));
      V = Q(:,:,n) + m * m';
      Ez(:,n) = [ m; 1 ];
      Ezz(:,:,n) = [ V, m; m', 1 ];
      
      % compute terms used to update W
      Wtop(:,I) = Wtop(:,I) + Ez(:,n) * x0';
      Wbottom(:,:,I) = bsxfun(@plus, Wbottom(:,:,I), Ezz(:,:,n));
    end
    clear x_n;

    %
    % M-STEP
    %  
    Waug = zeros( num_movies, D+1 );
    for m=1:num_movies
      Waug(m,:) = Wbottom(:,:,m) \ Wtop(:,m);
    end
    W_new = Waug( :, 1:(end-1) );
    xbar_new = Waug( :, end );    
    Ttot = zeros(num_movies, 1);
    for n=1:num_users
      x_n = S_train( :, n );    
      Tn = double(x_n ~= 0);
      Ttot = Ttot + Tn;

      predError = (x_n - Waug*Ez(:,n)).^2 + sum((W_new*Q(:,:,n)).*W_new, 2);
      Psi_new = Psi_new + Tn .* predError;
    end
    Psi_new = Psi_new ./ Ttot;
    clear x_n;

    % update parameters
    W = W_new;
    Psi = Psi_new;
    xbar = xbar_new;
    
    % compute complete data log-likelihood
    ll_this = 0;  H = 0;
    for n=1:num_users
      x_n = S_train( :, n );
      Tn = double(x_n ~= 0);

      ll_this = ll_this + sum(diag(Ezz(:,:,n)));
      H = H + D/2 * log( 2*pi*exp(1) ) + 0.5 * log( det( Q(:,:,n) ) );
      predError = (x_n - Waug*Ez(:,n)).^2 + sum((W_new*Q(:,:,n)).*W_new, 2);
      ll_this = ll_this + Tn'*(log(Psi) + predError./Psi);
    end
    clear x_n;
    loglike = -0.5 * ll_this + H;
    lls(iters) = loglike;
    diff_loglike = loglike - loglike_old;
    %fprintf('Delta Log-Like: %f\n', diff_loglike );

    % check for convergence
    if iters == max_iters
      done = true;
    elseif ( ( iters > 1 ) && ( diff_loglike <= conv_thresh ) )
      done = true;
    end

  end


  %
  % EVALUATE TEST DATA
  %
  R_train = zeros(size(S_train,1),D);
  for n=1:num_users
    x_n = S_train( :, n );    
    I = find( x_n );
    if isempty(I), continue; end; 

    % construct observed quantities
    W0 = W( I, : );
    invPsi0 = diag(1./Psi(I));
    x0 = x_n( I );
    xbar0 = xbar( I );

    % compute posterior mean
    G = inv( eye(D) + W0' * invPsi0 * W0 );
    Ez = G * W0' * invPsi0 * (x0 - xbar0);  

    % predict rankings
    R_train(n,:) = Ez';
  end

end
