% function [ mvn_model, BIC ] = train_gmm_bic( data, n_mvn_components, diag_only, verbose, lambda, max_iter, cv_thr )
%
% Wrapper around "do_em_gmm.m", 
% to initialize and train a GMM
% "mvn_model" is a structure array
%
% Note that we try several ways of initializing,
% and select the maximum BIC result (lambda = 1 by default).
%
% Only "data" and "n_mvn_components" are mandatory.
%   data is D x N where D is the dimensionality and N the number of data points.
%
% The remaining input arguments have the following default values:
%   diag_only = 0     full cov mat
%   verbose   = 1     set it to 2 if you want "do_em_gmm" to be verbose as well
%   lambda    = 1     used by the BIC model selection
%   max_iter  = 1000  used by do_em_gmm.m
%   cv_thr    = 1e-8  used by do_em_gmm.m
%
%
% Guillaume Lathoud
% 2006
% lathoud@idiap.ch

function [ mvn_model, BIC ] = train_gmm_bic( data, n_mvn_components, diag_only, verbose, lambda, max_iter, cv_thr )
  
  if nargin < 2
    error( [ mfilename '.train_gmm_bic() needs at least 2 input arguments.' ] );
  end

  if size( data, 2 ) < 2
    error( [ mfilename ' believes that fitting a GMM on a single data point is a foolish idea.' ] );
  end

  if n_mvn_components < 1
    error( [ mfilename ': n_mvn_components must be > 0.' ] );
  end
  
  if ~exist( 'diag_only', 'var' ),    diag_only = 0;    end
  if ~exist( 'verbose', 'var' ),      verbose = 1;      end
  if ~exist( 'lambda', 'var' ),       lambda = 1;       end
  if ~exist( 'max_iter', 'var' ),     max_iter = 1000;  end
  if ~exist( 'cv_thr', 'var' ),       cv_thr = 1e-8;    end
  

  % We fail, unless we succeed
  mvn_model = [];
  
  
  if n_mvn_components > 1
    
    % GMM with multiple componens: try multiple initializations, and pick the max BIC result
    
     % First, by splitting the data along the first component
    % ( e.g. C0 for MFCCs )
    
    clear init_mvn_model
    
    M = n_mvn_components;
    
    [ tmp, ind ] = sort( data( 1,: ) );
    data = data( :, ind );
    n = size( data, 2 );
    
    clear init_mvn_model
    
    for qqq = 1:M
      
      init_mvn_model( qqq ).w = 1 / M;
      
      s = 1 + round( (qqq-1)/M * ( n-1 ) );
      e = 1 + round( qqq/M * (n-1) );
      
      if e <= s + 1
	
	s = max( 1, s - 1 );
	e = min( n, e + 1 );
	
      end
      
      ind = (s:e);
      
      init_mvn_model( qqq ).mu = mean( data( :, ind ), 2 );
      
      init_mvn_model( qqq ).Sigma = cov( data( :, ind ).' );
      
      if ~all( size( init_mvn_model( qqq ).Sigma ) == [ size( data, 1 )  size( data, 1 ) ] )
	init_mvn_model( qqq ).Sigma = diag( ones( size( init_mvn_model( qqq ).mu ) ) );
      end
      
      if diag_only
	init_mvn_model( qqq ).Sigma = diag( diag( init_mvn_model( qqq ).Sigma ) );
	init_mvn_model( qqq ).diag_Sigma = 1;
      end
      
    end

    dbclear warning;
    warning off
    [ new_mvn_model_1, all_log_lkld_1 ] = do_em_gmm( init_mvn_model, data, max( 0, verbose-1 ), 0, max_iter, cv_thr );
    warning on
    dbstop warning;

    % Second, through random perturbation of the mean statistics
    
    mu    = mean( data, 2 );
    Sigma = cov( data.' );
    
    clear init_mvn_model

    for qqq = 1:M
      
      init_mvn_model( qqq ).w = 1 / M;

      init_mvn_model( qqq ).mu = mu .* ( 1 + 0.01 * normrnd( 0,1, size( mu,1 ), size( mu,2 ) ) );

      init_mvn_model( qqq ).Sigma = Sigma .* ( 1 + 0.01 * normrnd( 0, 1, size( Sigma, 1 ), size( Sigma, 2 ) ) );
      init_mvn_model( qqq ).Sigma = triu( init_mvn_model( qqq ).Sigma, 0 ) + triu( init_mvn_model( qqq ).Sigma, 1 ).';

      if ~all( size( init_mvn_model( qqq ).Sigma ) == [ size( data, 1 )  size( data, 1 ) ] )
	init_mvn_model( qqq ).Sigma = diag( ones( size( init_mvn_model( qqq ).mu ) ) );
      end
      
      if diag_only
	init_mvn_model( qqq ).Sigma = diag( diag( init_mvn_model( qqq ).Sigma ) );
	init_mvn_model( qqq ).diag_Sigma = 1;
      end
      
    end

    dbclear warning
    warning off
    [ new_mvn_model_2, all_log_lkld_2 ] = do_em_gmm( init_mvn_model, data, max( 0, verbose-1 ), 0, max_iter, cv_thr );
    warning on
    dbstop warning

    % Third, through yet another deterministic initialization

    init_mvn_model = determinit_gmm( M, data );
    
    if diag_only
      for qqq = 1:numel( init_mvn_model )
	init_mvn_model( qqq ).Sigma = diag( diag( init_mvn_model( qqq ).Sigma ) );
	init_mvn_model( qqq ).diag_Sigma = 1;
      end
    end

    dbclear warning
    warning off
    [ new_mvn_model_3, all_log_lkld_3 ] = do_em_gmm( init_mvn_model, data, max( 0, verbose-1 ), 0, max_iter, cv_thr );
    warning on
    dbstop warning

    %%%

    % Estimate the number of free parameters of each model

    K_1 = estimate_K_gmm( new_mvn_model_1 );
    K_2 = estimate_K_gmm( new_mvn_model_2 );
    K_3 = estimate_K_gmm( new_mvn_model_3 );
    
    % Compute the BIC score of each resut
    
    N = size( data, 2 );
    
    BIC_1 = all_log_lkld_1 - lambda / 2 * K_1 * log( N );
    BIC_2 = all_log_lkld_2 - lambda / 2 * K_2 * log( N );
    BIC_3 = all_log_lkld_3 - lambda / 2 * K_3 * log( N );
    
    % Pick the one with maximum BIC score

    [ max_BIC, max_ind ] = max( [ BIC_1  BIC_2  BIC_3 ] );

    switch max_ind
     case 1
      mvn_model = new_mvn_model_1;
     case 2
      mvn_model = new_mvn_model_2;
     case 3
      mvn_model = new_mvn_model_3;
    end
    
    if verbose
      disp( sprintf( [ mfilename '.train_gmm_bic():  fitted model -> BIC_1: %.10g,  BIC_2: %.10g,  BIC_3: %.10g,  max_ind: %d' ], BIC_1, BIC_2, BIC_3, max_ind ) );
    end

  end
    
  
  
  % We can always fall back on the single Gaussian
  
  if isempty( mvn_model )
    
    if verbose & ( n_mvn_components > 1 )
      disp( sprintf( [ mfilename '.train_gmm_bic(): training failed, falling back on single Gaussian.' ] ) );
    end

    [ mu, Sigma ] = mean_cov( data, [] );
    
    mvn_model.w = 1;
    
    mvn_model.mu = mu;
    
    [ R, p ] = chol( Sigma );
    
    if ( p == 0 ) & ~diag_only
      
      mvn_model.Sigma = Sigma;
      
    else	
      
      mvn_model.Sigma      = diag( diag( Sigma ) );
      mvn_model.diag_Sigma = 1;
      
      % Prevent a really rare, extremely bad case
      % mainly due to pervert users!
      
      [ R, p ] = chol( mvn_model.Sigma );
      
      if p ~= 0
	
	if verbose
	  disp( [ mfilename '.train_gmm_bic(): could not fit even a single diagonal Gaussian, falling back on Id covariance matrix.' ] );
	end
	
	mvn_model.Sigma = eye( numel( mvn_model.mu ) );
	
      end
      
    end
    
    
  end % if training failed
  
    
  % Finally mark which ones have diagonal covariance matrix

  for m = 1:numel( mvn_model )
    
    mvn_model( m ).diag_Sigma = all( all( mvn_model( m ).Sigma == diag( diag( mvn_model( m ).Sigma ) ) ) );
    
    mvn_model( m ).w          = full( mvn_model( m ).w );
    mvn_model( m ).mu         = full( mvn_model( m ).mu );
    mvn_model( m ).Sigma      = full( mvn_model( m ).Sigma );
    
  end

  if verbose
    disp( sprintf( [ mfilename ': %d components were asked for. I am returning a GMM with %d components (%d/%d are diagonal)' ], n_mvn_components, numel( mvn_model ), sum( [ mvn_model.diag_Sigma ] ), numel( mvn_model ) ) );
  end

  
  % If the user wants the BIC score, then let's return it
  if nargout > 1
    
    [ tmp, tmp, tmp, log_lkld ] = compute_post_gmm( mvn_model, data );
    
    K   = estimate_K_gmm( mvn_model );
    
    N   = size( data, 2 );
    
    BIC = log_lkld - lambda * K / 2 * log( N );
    
  end
  

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [a_mean, a_cov, diag_cov] = mean_cov( data, weight )
  
  if nargin < 1
    disp( [ mfilename '.mean_cov() needs at least 1 input argument.' ] );
  end

  if ~exist( 'weight', 'var' )
    weight = [];
  end

  if isempty( weight )
    
    a_mean = mean( data.' ).';
    
    % Full covariance matrix
    a_cov = cov( data.' );
    
  else
    
    sum_weight = max( sum( weight ), eps );

    a_mean   = data * weight(:) / sum_weight;
    
    tmp_term = full( data * diag( sparse( sqrt( weight / sum_weight ) ) ) );
    
    % Full covariance matrix
    a_cov = tmp_term * tmp_term.';
    
  end

  
  % Make sure that it is positive definite ( for the inversion in log_mvnpdf.m )
  [ R, p ] = chol( a_cov );
  
  diag_cov = ( p~= 0 );
  
  if diag_cov
    % If not positive definite, fall back on the diagonal covariance matrix
    a_cov = diag( var( data.' ) );
  end

  a_mean = full( a_mean );
  a_cov  = full( a_cov );
