%************************************************************
% File:         mixtureModel_online.m
% Date:         September 4, 2006
% Author:       Jo-Anne Ting
% Description:
%  Real-time (online) implementation of Mixture Model 
%  for globally linear regression (assume data is generated 
%  from a Gaussian distribution and outliers are generated 
%  from a uniform distribution).
%
%  To simulate a uniform distribution, we use a Gaussian
%  distribution with a very wide, fixed variance.
% 
%  A forgetting factor, lambda, is used for real-time
%  implementation.
%
% Inputs:
%
%  (where N = number of data samples,
%         d = dimensionality of input data)
%
%  X               : training input (Nxd matrix)
%  Y               : training output (Nx1 matrix)
%  n_iter          : max number of EM iterations
%  tol             : tolerance for termination of
%                    algorithm (proportional change in
%                    likelihood)
%  lambda          : forgetting factor, [0, 1]
%  ss              : sufficient statistics collected so far
%  nn              : nn-th data sample observed so far
%
% Outputs:
%
% A data structure "s" containing the following:
%
%  s.alpha1       : probability that samples are inliers
%  s.alpha2       : probability that samples are outliers
%  s.mu1          : mean of "inlier" Gaussian component
%                   (capturing inliers)
%  s.Sigma1       : covariance of "inlier" Gaussian component
%  s.mu2          : mean of "outlier" Gaussian component
%                   (capturing outliers)
%  s.Sigma2       : covarianace of outlier Gaussian 
%                   component  
%  s.z1           : membership weights for data cloud
%  s.z2           : membership weights for outlier component
%  s.beta_sigma2  : variance of beta
%  s.beta         : regression coefficietn estimate
%
%  as well as other sufficient statistics:
%   s.sum_z1yx  
%   s.sum_z1xx  
%   s.sum_z1xxT 
%   s.sum_z1
%   s.sum_z2 
%   s.sum_z1x
%   s.sum_z2x
%   s.sum_zxu2  
%   s.n 
%
%************************************************************
function [s] = mixtureModel_online(...
            X, Y, n_iter, tol, lambda, ss, nn)

% Pad input vector with an offset dimension
X1  = [X Y];
Xpadded = [X 1];
[N,d] = size(X1);

% Setup of priors
%-------------------------------------------------------
s.alpha1 = 0.5;
s.alpha2 = 1-s.alpha1;
s.mu1    = zeros(d,1);
s.Sigma1 = 1e3*eye(d);
s.mu2    = zeros(d,1);
% Fix to be very large (covers noise)
s.Sigma2 = 1e3*eye(d);   
% To avoid computational problems
SMALL    = 1e-4;

% For first data sample, store sufficient statistics
%-------------------------------------------------------
if nn == 1
  % Collect sufficient statistics
  s.sum_z1yx  = Xpadded'*Y;
  s.sum_z1xx  = Xpadded'*Xpadded;
  s.sum_z1xxT = X1'*X1;
  s.sum_z1x   = X1;
  s.sum_z1    = 1;
  s.sum_z2x   = zeros(N,d);
  s.sum_z2    = 0;
  s.sum_zxu2  = (X1 - s.mu1')'*(X1 - s.mu1');
  s.n         = 1;

  % Assume this point is NOT an outlier
  s.z1        = 1;
  s.z2        = 0;
  s.beta_sigma2 = pinv(s.sum_z1xx);
  s.beta      = pinv(s.sum_z1xx)*s.sum_z1yx;
  return;
end

% For other data points, iterate on each data point
%-------------------------------------------------------
for i=1:n_iter,

  % E-step
  %-----------------------------------------------------
  % Calculate membership weights (nn-th term) 
  z1_num = s.alpha1 / (sqrt((2*pi)^d) * ...
     sqrt(det(s.Sigma1 + SMALL*eye(d)))).* ...
      exp( -0.5*(X1 - s.mu1') * ...
          inv(s.Sigma1 + SMALL*eye(d))*...
          (X1 - s.mu1')'...
         );
  z2_num = s.alpha2 / (sqrt((2*pi)^d) * ...
     sqrt(det(s.Sigma2))) .* ...
     exp( -0.5*(X1 - s.mu2') * ...
           inv(s.Sigma2)*...
            (X1 - s.mu2')'...
         );
  s.z1 = z1_num ./ (z1_num + z2_num);
  s.z2 = z2_num ./ (z1_num + z2_num);
 
  % Variance of z1
  s.z1_sigma = s.z1*(1-s.z1);
 
  % M-step
  %-----------------------------------------------------
  % Collect sufficient statistics
  sum_z1    = s.z1 + lambda*ss.sum_z1;
  sum_z2    = s.z2 + lambda*ss.sum_z2;
  sum_z1x   = X1.*s.z1 + lambda*ss.sum_z1x;
  sum_z2x   = X1.*s.z2 + lambda*ss.sum_z2x;
  sum_z1xxT = (s.z1).*(X1'*X1) + lambda*ss.sum_z1xxT;
  sum_zxu2  = s.z1.*((X1 - s.mu1')'*(X1 - s.mu1')) ...
              + lambda*ss.sum_zxu2;

  s.alpha1 = sum_z1 / (1 + lambda*ss.n); 
  s.alpha2 = sum_z2 / (1 + lambda*ss.n); 

  s.Sigma1 = sum_zxu2 ./ sum_z1;
  s.mu1 = (sum_z1x ./ sum_z1)';  % dx1
  s.mu2 = (sum_z2x ./ (sum_z2+SMALL))';  % dx1
end

% Store sufficient statistics
%-------------------------------------------------------
s.sum_z1yx  = s.z1*(Xpadded'*Y) + lambda*ss.sum_z1yx;
s.sum_z1xx  = s.z1.* (Xpadded'*Xpadded) + ...
              lambda*ss.sum_z1xx;
s.sum_z1xxT = s.z1.* (X1'*X1) + lambda*ss.sum_z1xxT;
s.sum_z1    = s.z1 + lambda*ss.sum_z1;
s.sum_z2    = s.z2 + lambda*ss.sum_z2;
s.sum_z1x   = X1.*s.z1 + lambda*ss.sum_z1x;
s.sum_z2x   = X1.*s.z2 + lambda*ss.sum_z2x;
s.sum_zxu2  = s.z1.*((X1 - s.mu1')'*(X1 - s.mu1')) + ...
              lambda*ss.sum_zxu2;
s.n         = 1 + lambda*ss.n;

% Calculate regression estimate
%-------------------------------------------------------
s.beta_sigma2 = inv(s.sum_z1xx+1e-3);
s.beta = s.beta_sigma2*s.sum_z1yx;
