function res = Gauss(mode, varargin)
  % usage is one of the two following:
  %    model = Gauss('train', X, Y, covarianceType, sharedCovariance);
  %    Y     = Gauss('predict', model, X);
  %
  % if covarianceType == 0, we just use identity; if covarianceType == 1,
  % we use diagonal covariance; if covarianceType == 2, we use full
  % covariance.  sharedCovariance should be 0 or 1, dictating whether
  % all classes get the same covariance or not.

  switch lower(mode),
   case 'name'
    res = 'Gauss';
    
   case 'train'
    if nargin ~= 5,
      error('Gauss training requires four arguments: X, Y, covarianceType and sharedCovariance');
    end
    % ma
    [NX,DX] = size(varargin{1});
    [NY,DY] = size(varargin{2});
    if NX ~= NY,
      error('there must be the same number of data points in X and Y');
    end;
    if DY ~= 1,
      error('Y must have only one column');
    end
    if ~isscalar(varargin{3}),
      error('covarianceType must be a scalar');
    end;
    if ~(varargin{3} == 0 || varargin{3} == 1 || varargin{3} == 2),
      error('covarianceType must be 0, 1 or 2');
    end;
    if ~isscalar(varargin{4}),
      error('sharedCovariance must be a scalar');
    end;
    if ~(varargin{4} == 0 || varargin{4} == 1),
      error('sharedCovariance must be 0 or 1');
    end;
    
    % return "model" -- for DT, we have to build it
    [res.classPi, res.mu, res.si2] = GaussTrain(varargin{1}, varargin{2}, varargin{3}, varargin{4});
    
    
   case 'predict'
    % check arguments
    if nargin ~= 3,
      error('prediction requires two arguments: the model, and X');
    end;
    model = varargin{1};
    X = varargin{2};
    
    if ~isfield(model,'classPi'),
      error('model does not appear to be a Gauss model');
    end;
    if ~isfield(model,'mu'),
      error('model does not appear to be a Gauss model');
    end;
    if ~isfield(model,'si2'),
      error('model does not appear to be a Gauss model');
    end;
    
    % set up output
    N = size(X,1);
    res = zeros(N,1);
    
    % compute predictions
    res = GaussPredict(model.classPi, model.mu, model.si2, X);
  
   otherwise
    error('unknown Gauss mode: need "train" or "predict"');
  end;

  
  
function [classPi, mu, si2] = GaussTrain(X, Y, covarianceType, sharedCovariance)
  % we want to train a generative Gaussian model.  if there are K labels, then
  % classPi(k) should be the prior probability that we choose label (k).  each
  % class is then represented by a Gaussian with mean mu{k} and covariance
  % si2{k}.  mu{k} should *always* be a 1xD vector and si2{k} should *always* 
  % be a DxD matrix.  for simplicity, we may assume k = 1....K
  %
  % if sharedCovariance==1 then all Gaussians should use the same
  % covariance.  that is, si2{k} == si2{l} for all pairs k/l.  (yes, you
  % should store K copies of it even though this is technically wasteful.
  % it will just make prediction a bit easier.)  
  %  
  % if covarianceType==0, then all covariances should just be the identity 
  % matrix, eye(D).
  % 
  % if covarianceType==1, then the covariances should all be diagonal matrices
  % (i.e., the off-diagonal elements should be zero).  this says that features
  % are essentially independent.
  %  
  % if covarianceType==2, then the covariances should be full DxD matrices.
  
  % find dimensions and # of classes
  [N,D] = size(X);
  K = max(Y);
  
  % initialize variables
  classPi = zeros(1,K);
  mu = {};
  si2 = {};
  for k=1:K,
    mu{k} = zeros(1,D);
    si2{k} = eye(D);    % note we initialize to identity
  end;
    
  % first, compute the prior probabilities of each class
  for k=1:K,
    classPi(k) = length(find(Y == k));
  end;
  % normalize
  classPi = classPi/N;
  
  % next, compute the means of the Gaussians for each class
  for k=1:K,
      if length(find(Y == k)) == 1
          mu{k} = X((Y == k),:);
      elseif length(find(Y == k)) == 0 
          mu{k} = [];
      else
          mu{k} = sum(X((Y == k),:))/length(find(Y == k));
      end;
  end;
  
  % finally compute the variances.  first check to see if we're going to
  % compute one overall covariance or per-class covariances.
  if covarianceType ~= 0,   % the "0" case is done for you already!
    if sharedCovariance,
      % compute a single covariance
      bigSI2 = zeros(D,D);
      for k = 1:K          
          X1 = X(Y == k,:);     % Pick up the data in k class.
          % It's possible that X1 is a empty matrix if no data in class k.
          % If that's the case, x_mean is also empty and nothing will be
          % added to bigSI2.So we dont care it.

          for d = 1:D
              if size(mu{k}, 1) > 0 
                    X1(:,d) = X1(:,d) - mu{k}(d);   % remove the mean for k class   
              end;
          end;
          
          bigSI2 = bigSI2 + X1' * X1;   %variance for all data
      end;
      bigSI2 = bigSI2/(N);  % Use N instead of N-1, as Hal seems to use N.
          

      % flip it to diagonal if we need to
      if covarianceType == 1,
        bigSI2 = diag(diag(bigSI2));
      end;
      
      % copy it to all the si2s
      for k=1:K,
        si2{k} = bigSI2;
      end;
    else
      % compute a different covariance for each class
      for k=1:K,
        % compute a big covariance for class k
        N1 = size(X(Y == K,:), 1);  % # of data points for this class k
        si2{k} = cov(X(Y == k,:)) * (N1 - 1) / N1;
        
        % diagonalize it if necessary
        if covarianceType == 1,
          si2{k} = diag(diag(si2{k}));
        end;
      end;
    end;
  end;
  
  % we're done!
  
  
function y = GaussPredict(classPi, mu, si2, X)
  % classPi, mu, si2 are just like in GaussTrain; X is an NxD vector

  % get dimensions
  [N,D] = size(X);
  K = length(classPi);
  y = zeros(N,1);

  % for each class k, compute lp(n,k) = log probability of class k given X(n,:)
  lp = zeros(N,K);
  for k=1:K,
    % precompute inverse and determinant:
    % for high-D, add a small identity to avoid singularity
    if D<=2, 
      invSI2 = inv(si2{k});
      logdetSI2 = logdet(si2{k});
    else
      invSI2 = inv(si2{k} + 0.1 * eye(D));
      logdetSI2 = logdet(si2{k} + 0.1 * eye(D));    
    end;
    
    % compute log probability for each point:
    %
    % remember p(class=k | x_n) = (1/Z) p(class=k) p(x_n | class=k)
    %                           = (1/Z) Mult(k | classPi) Nor(x_n | mu_k si2_k)
    %
    % and Nor(x | mu, si2) = 1/sqrt(2*pi*det(si2)) exp[-(1/2) (x-mu) inv(si2) (x-mu)']
    %
    % we have precomputed inv(si2) and [log det(si2)] for you above
    for n=1:N,
        if size(mu{k}, 1) ~= 0
                lp(n,k) = log(classPi(k)) - 0.5 * (log(2*pi) + logdetSI2) - 0.5 * (X(n,:)-mu{k}) * invSI2 * (X(n,:)-mu{k})';
        end;

    end;
  end;
  
  % compute the most likely class y
  for n=1:N,
    [v y(n)] = max(lp(n,:));
  end;
  
  
function y = logdet(A)
  U = chol(A);
  y = 2*sum(log(diag(U)));

  