function [ new_mvn_model, all_log_lkld ] = do_em_gmm( mvn_model, data, verbose, log_domain_flag, max_iter, cv_prop )

% function [ new_mvn_model, all_log_lkld ] = do_em_gmm( mvn_model, data, [ verbose ], [ log_domain_flag ], [ max_iter ], [ cv_prop ] )
%  
% EM optimization of a GMM (full covariance matrix).
% The GMM has M components mvn_model( 1:M ).
%
% By default log_domain_flag = 0 (linear domain is the default),
% because so far the log domain implementation does not bring much.
%
% See also: estimate_K_gmm.m, compute_post_gmm.m.
%
% By Guillaume Lathoud, 2006.
% lathoud@idiap.ch

  
  if nargin < 2
    error( [ mfilename '.do_em_mvn needs at least 2 input arguments.' ] );
  end

  % Default: verbosity active
  if ~exist( 'verbose', 'var' )
    verbose = [];
  end
  
  if isempty( verbose )
    verbose = 1;
  end
  
  % M-step maybe implemented in log domain (maybe more precise, but cumbersome)
  if ~exist( 'log_domain_flag', 'var' )
    log_domain_flag = [];
  end
  
  if isempty( log_domain_flag ) 
    log_domain_flag = 0;
  end
    
  % Default: max_iter = 1000
  
  if ~exist( 'max_iter', 'var' )
    max_iter = [];
  end
  
  if isempty( max_iter )
    max_iter = 1000;
  end
  
  % Default proportion for the convergence test: 1e-8
  
  if ~exist( 'cv_prop', 'var' )
    cv_prop = [];
  end
  
  if isempty( cv_prop )
    cv_prop = 1e-8;
  end
  
  % Dimensionality
  
  D = size( data, 1 );
  
  % Number of samples
  
  N = size( data, 2 );
  
  % Number of components,
  % each component is a multivariate normal distribution
  
  M = numel( mvn_model );

  % Binary flag to mark which covariance matrix has turn'd diagonal
  if ~isfield( mvn_model, 'diag_Sigma' )
    mvn_model( 1 ).diag_Sigma = [];
  end

  for m = 1:M
    if isempty( mvn_model( m ).diag_Sigma )
      % By default, full covariance matrix
      mvn_model( m ).diag_Sigma = 0;
    end
  end
  
  % EM
  finished = 0;
  
  iter = 0;
  previous_all_log_lkld = -Inf;
  
  while ~finished

    if verbose
      disp( ' ' );
    end
    
    % First make sure that all covariance matrices are positive definite
    % If not, fall back on a diagonal covariance matrix (change of model)
    % for ONE particular matrix.
    % -> this prevents many crashes.
    for m = 1:M
      
      if ~mvn_model( m ).diag_Sigma
	
	[ R, p ] = chol( mvn_model( m ).Sigma );
	mvn_model( m ).diag_Sigma = ( p ~= 0 );
	
	if mvn_model( m ).diag_Sigma
	  
	  if verbose
	    disp( sprintf( [ mfilename ': from now on, forcing covariance matrix of model %d to be diagonal.' ], m ) );
	  end
												    
	  % Then we need to compute the lkld again
	  % ( the convergence test won't be valid )
	  previous_all_log_lkld = -Inf;
	  all_log_lkld          = -Inf;
	  
	end
	
      end
      
      if mvn_model( m ).diag_Sigma
	
	% use diagonal only
	mvn_model( m ).Sigma = diag( sparse( diag( mvn_model( m ).Sigma ) ) );
	
      end
      
    end

    tmp = find( [ mvn_model.diag_Sigma ] );
    
    if verbose
      if ~isempty( tmp )
	disp( [ mfilename ': E-step: using diagonal only for models [' sprintf( ' %d', tmp ) ' ]' ] );
      end
    end
    
    % E step
    
    % Compute likelihoods
    %
    % We use a trick to avoid variance flooring
    % -> which means we may need to drop a component
    % and restart EM using the present state of the remaining
    % parameters
    
    is_stable = 0;

    while ~is_stable
      
      M = numel( mvn_model );
      
      component_to_keep = ones( 1, M );
      
      comp_log_joint    = repmat( NaN, M, N );
      
      
      ttt = dbstatus;
      
      warn_stop = ismember( 'warning', { ttt.cond } );
      dbclear warning;

      error_stop = ismember( 'error', { ttt.cond } );
      dbclear error;
      
      for m = 1:M
	
	try
	
	  lastwarn( '' );
	  
	  % log_mvnpdf may throw a warning or an error
	  comp_log_joint( m,: ) = log( mvn_model( m ).w ) + reshape( log_mvnpdf( data.', mvn_model( m ).mu(:).', full( mvn_model( m ).Sigma ) ), 1, [] );
	  
	  % Detect a possible unstability (if log_mvnpdf threw a warning)
	  component_to_keep( m ) = isempty( lastwarn );
	  
	  % Detect another possible unstability (NaN in the result)
	  component_to_keep( m ) = component_to_keep( m ) & ~any( isnan( comp_log_joint( m,: ) ) );

	  % Detect another possible unstability (+Inf in the result)
	  % It can happen if Sigma is collapsing, but still well-conditioned
	  % ( likelihood around the mean is beyond machine's precision )
	  component_to_keep( m ) = component_to_keep( m ) & ~any( +Inf == comp_log_joint( m,: ) );
	  	  
	catch
	
	  % Detect a possible unstability (if log_mvnpdf threw an error)
	  component_to_keep( m ) = 0;
	  
	end
	
      end
      
      if warn_stop,	dbstop warning; end;
      if error_stop,    dbstop error; end;

      
      is_stable = all( component_to_keep );

      if ~is_stable

	% Drop unstable components
	
	if verbose
	  disp( sprintf( [ mfilename ': %d unstable components, we''re going to drop them.' ], sum( double( ~component_to_keep ) ) ) );
	end
	
	mvn_model = mvn_model( find( component_to_keep ) );
	
	% if isempty

	if isempty( mvn_model )
	  disp( [ mfilename ': lost all components! Returning empty.'  ] );
	  
	  new_mvn_model = [];
	  all_log_lkld = -Inf;
	  return
	end

	M = numel( mvn_model );
	
	% Rescale priors
	sum_w = sum( [ mvn_model.w ] );
	for m = 1:M
	  mvn_model( m ).w = mvn_model( m ).w / sum_w;
	end
	
	% The parameter space has changed -> restart EM
	previous_all_log_lkld = -Inf;
	all_log_lkld          = -Inf;
	
      end

    end

    % Compute joints and posteriors
    
    if M > 1
      log_lkld = my_logsum_fast( comp_log_joint );
    else
      log_lkld = comp_log_joint;
    end
    
    all_log_lkld = sum( log_lkld );
    
    comp_log_post = comp_log_joint - repmat( log_lkld, M, 1 );
    
    % M step

    if log_domain_flag

      % DEBUG

      if verbose
	disp( [ mfilename ': M-step, log domain' ] )
      end
    
      % M step: implementation in log domain
      % ( more precise but slower than linear domain )

      new_log_w = my_logsum_fast( comp_log_post.' );
      new_log_w = new_log_w - my_logsum_fast( new_log_w );
      
      new_w = exp( new_log_w );
      new_w = new_w / sum( new_w );
      
      for m = 1:M
	
	% Update weight of this component
	
	mvn_model( m ).w = new_w( m );
	
	% Update parameters of this component
	
	% Weights (posteriors scaled so that they sum to one over all samples).
	tmp_log_w     = comp_log_post( m,: ) - my_logsum_fast( comp_log_post( m,: ) );

	% Mean
	
	for i = 1:D

	  % Sum positive terms

	  ind = find( data( i,: ) > 0 );
	  tmp = exp( my_logsum_fast( log( data( i, ind ) ) + tmp_log_w( ind ) ) );

	  % Sum negative terms

	  ind = find( data( i,: ) < 0 );
	  tmp = tmp - exp( my_logsum_fast( log( -data( i, ind ) ) + tmp_log_w( ind ) ) );
	  
	  % Result

	  mvn_model( m ).mu( i ) = tmp;

	end

	% Covariance matrix
	% ( use of log avoids to reach the limits of the machine in most cases )

	tmp_term = data - repmat( mvn_model( m ).mu(:), 1, N );
	
	mvn_model( m ).Sigma = zeros( D, D );

	for i = 1:D

	  j_max = D;
	  
	  if mvn_model( m ).diag_Sigma
	    j_max = i;
	  end
	  
	  for j = i:j_max

	    tmp_product = tmp_term( i,: ) .* tmp_term( j,: );
	    
	    % Sum positive terms
	    
	    ind = find( tmp_product > 0 );
	    tmp = exp( my_logsum_fast( log( tmp_product( ind ) ) + tmp_log_w( ind ) ) );
	    
	    % Sum negative terms
	    
	    ind = find( tmp_product < 0 );
	    tmp = tmp - exp( my_logsum_fast( log( -tmp_product( ind ) ) + tmp_log_w( ind ) ) );
	    
	    % Store it
	    
	    mvn_model( m ).Sigma( i, j ) = tmp;
	    
	  end

	end
	
	% Covariance matrix is symmetric

	mvn_model( m ).Sigma = mvn_model( m ).Sigma + triu( mvn_model( m ).Sigma, 1 ).';

      end
      
    else

      if verbose
	% DEBUG
	disp( [ mfilename ': M-step, linear domain'  ] );
      end

      % M step: implementation in linear domain
      % ( less precise but faster than log domain )

      new_w = exp( my_logsum_fast( comp_log_post.' ) );
      new_w = new_w / sum( new_w );
      
      for m = 1:M
	
	% Update weight of this component
	
	mvn_model( m ).w = new_w( m );
	
	% Update parameters of this component
	
	% Weights
	tmp_w     = exp( comp_log_post( m,: ) );
	tmp_sum_w = sum( tmp_w );
	
	% Mean
	
	mvn_model( m ).mu = full( ( data * tmp_w(:) ) / tmp_sum_w );

	% Covariance matrix

	
	
	
	if 0

	  tmp_term = full( ( data - repmat( mvn_model( m ).mu(:), 1, N ) ) * diag( sparse( sqrt( tmp_w(:) / tmp_sum_w ) ) ) );
	  
	  mvn_model( m ).Sigma = zeros( D, D );
	  
	  for i = 1:D
	    
	    j_max = D;
	    
	    if mvn_model( m ).diag_Sigma
	      j_max = i;
	    end
	    
	    for j = i:j_max
	      
	      mvn_model( m ).Sigma( i, j ) = tmp_term( i,: ) * tmp_term( j,: ).';
	      
	    end
	    
	  end

	  % Covariance matrix is symmetric
	  
	  mvn_model( m ).Sigma = mvn_model( m ).Sigma + triu( mvn_model( m ).Sigma, 1 ).';

	else

	  
	  if ~mvn_model( m ).diag_Sigma 
	    
	    % Update full covariance matrix
	    tmp_term = full( ( data - repmat( mvn_model( m ).mu(:), 1, N ) ) * diag( sparse( sqrt( tmp_w(:) / tmp_sum_w ) ) ) );
	    mvn_model( m ).Sigma = tmp_term * tmp_term.';
	    
	  else
	    
	    % Update diagonal covariance matrix
	    tmp_var = ( ( data - repmat( mvn_model( m ).mu(:), 1, N ) ) .^ 2 ) * tmp_w(:) / tmp_sum_w;
	    mvn_model( m ).Sigma = diag( sparse( tmp_var ) );
	    
	  end

	end
	
      end
      

    end



    for m = 1:M

      % If Sigma not definite positive,
      % then restrict to the diagonal
      
      if ~mvn_model( m ).diag_Sigma

	[R,p] = chol( mvn_model( m ).Sigma );
	
	mvn_model( m ).diag_Sigma = ( p~=0 );

	if mvn_model( m ).diag_Sigma

	  if verbose
	    disp( sprintf( [ mfilename ': from now on, forcing covariance matrix of model %d to be diagonal.' ], m ) );
	  end
												    
	  % Then we need to compute the lkld again
	  % ( the convergence test won't be valid )
	  previous_all_log_lkld = -Inf;
	  all_log_lkld          = -Inf;
	  
	end
	
      end
      
      if mvn_model( m ).diag_Sigma
	
        % use diagonal only
        mvn_model( m ).Sigma = diag( sparse( diag( mvn_model( m ).Sigma ) ) );
	
      end
      
    end
    
    if verbose
      disp( [ mfilename sprintf( ': iter:%d   all_lkld:%.10g', iter, all_log_lkld ) ] );
    end
    
    
    % Termination test
    
    iter = iter + 1;
    
    % Note that when both "previous_all_log_lkld" and "all_log_lkld" are -Inf,
    % we have -Inf - -Inf = NaN, NaN < abs(...) * 1e-8 is always 0,
    % so that we basically don't check convergence, which is what we need
    % in such case ( change of parameter space = restart EM ).
    
    finished = ( iter >= max_iter ) | ( ( all_log_lkld - previous_all_log_lkld ) < abs( all_log_lkld ) * cv_prop );
    
    previous_all_log_lkld = all_log_lkld;
    
  end
  
  % Final check: make sure all matrices are positive definite
  % and well-conditionned
  
  is_okay = logical( ones( numel( mvn_model ), 1 ) );
  
  for m = 1:numel( mvn_model )

    mvn_model( m ).Sigma = full( mvn_model( m ).Sigma );
    
    [ R, p ] = chol( mvn_model( m ).Sigma );
    
    is_okay( m ) = ( p == 0 ) & ( rcond( mvn_model( m ).Sigma ) > 1000 * eps );
    
  end
  
  mvn_model = mvn_model( find( is_okay ) );
  
  % Update weights if necessary
  
  if any( ~is_okay ) & ( ~isempty( mvn_model ) )
    
    w = [ mvn_model.w ];
    w = w / sum( w );
    
    for m = 1:numel( mvn_model )
      mvn_model( m ).w = w( m );
    end

  end

  new_mvn_model = mvn_model;

  % Also return the lkld

  if isempty( mvn_model )

    % In case of failure

    all_log_lkld = -Inf;

  else

    % In case of success

    [ tmp, tmp, tmp, all_log_lkld ] = compute_post_gmm( mvn_model, data );

  end
  
