%************************************************************
% File:         sparse_mixture_online.m
% Date:         January 24, 2007
% Author:       Jo-Anne Ting
% Description:
%  Implements online/realt-time version of Faul & Tipping's
%  Variational Bayesian Robust regression (2001) using a
%  forgetting/decay factor, lambda.
%
% 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 final converged
% values of:
%
%  s.z_a          : probability that samples were generated
%                   by a Gaussian distribution
%  s.z_b          : probability that samples were generated
%                   by a uniform distibution
%  s.z            : [0, 1]; indicates prob. that samples
%                   were generated by the Gaussian component
%  s.beta_sigma2  : covariance of beta
%  s.beta         : regression estimate
%  s.alpha        : hyperparameters for ARD on beta 
%  s.theta        : probability that sample is not an outlier
%  s.b1           : precision of Gaussian component
%
% as well as other sufficient statistics.
%
%************************************************************
function [s] = sparse_mixture_online(...
                X, Y, n_iter, tol, lambda, ss, nn)

% Augment input data
X1    = [X 1];
[N,d] = size(X1);

% Setup of priors
%-------------------------------------------------------
s.beta_sigma2 = 1e6*eye(d);
s.beta = zeros(d,1);
s.b1 = ss.true_varY;   
s.b2 = s.b1/2;    
s.alpha = s.b1*d*ones(d,1);     
s.theta = 0.5;
% To avoid computational problems
SMALL = 1e-6;


% For first data point, collect sufficient statistics
%-------------------------------------------------------
if nn == 1
  s.beta_sigma2 = 1e6*eye(d);
  s.beta = zeros(d,1);
  
  % Assume this point is NOT an outlier
  % Collect sufficient statistics
  s.maxY = Y;
  s.minY = Y;
  s.z = 1;
  s.sum_z = s.z;
  s.sum_n = 1;
  s.beta  = pinv(X1'*X1)*X1'*Y;
  s.sum_za = s.z;
  s.varY = 0;
  s.sumY = Y;
  s.sum_Ym2 = Y^2;
  s.sum_zxx = X1'*X1;
  s.sum_zxy = X1'*Y;
  return;
end

% Update sufficient statistics (for other samples)
%-------------------------------------------------------
maxY = ss.maxY;
minY = ss.minY;
% Update the max & min Y values seen so far 
if Y > ss.maxY
  maxY = Y;
end
if Y < ss.minY
  minY = Y;
end
s.sumY = Y + ss.sumY;
meanY = s.sumY / (1 + ss.sum_n);
s.sum_Ym2 = (Y - meanY)^2 + lambda*ss.sum_Ym2;
s.varY = s.sum_Ym2 / (1 + lambda*ss.sum_n);

% EM-Iterations
%-------------------------------------------------------
for i=1:n_iter,

  % E-step 
  %-----------------------------------------------------
  % Iterate through the E-step couple of times 
  % (suggested by Tipping)
 
  for nn=1:5

  % Calculate hidden variables
  %-----------------------------------------------------
  s.a = (Y - X1*s.beta).^2 + X1*s.beta_sigma2*X1';
  s.z_a = s.theta*sqrt(s.b1./(2*pi))*exp(-s.b1*s.a*0.5);
  %s.z_b = (1-s.theta)/(maxY - minY + SMALL);
  s.z_b = (1-s.theta)/1e4;
  s.z = s.z_a ./ (s.z_a + s.z_b);
 
  % Re-estimate the regression coefficients 
  %-----------------------------------------------------
  % Update sufficient statistics
  sum_zxx = (X1.*repmat(s.z,1,d))'*X1 + lambda*ss.sum_zxx;
  sum_zxy = (X1.*repmat(s.z,1,d))'*Y + lambda*ss.sum_zxy;
 
  % Assume outlier comes from a uniform distribution
  s.beta_sigma2 = inv(SMALL + s.b1*sum_zxx + diag(s.alpha) );
  s.beta = s.beta_sigma2*(s.b1*sum_zxy);
  end

  % M-step 
  %-----------------------------------------------------
  % Update sufficient statistics
  sum_z = s.z + lambda*ss.sum_z;
  sum_za = s.z.*s.a + lambda*ss.sum_za;
  sum_n = 1 + lambda*ss.sum_n;

  s.gamma = 1 - s.alpha.*diag(s.beta_sigma2);
  s.alpha = s.gamma ./ (s.beta.^2 + SMALL);
  s.theta = sum_z / sum_n;
  s.b1 = sum_z / (sum_za + SMALL);
  
end

% Store sufficient statistics
%-------------------------------------------------------
s.maxY = maxY;
s.minY = minY;
s.sum_zxx = (X1.*repmat(s.z,1,d))'*X1 + lambda*ss.sum_zxx;
s.sum_zxy = (X1.*repmat(s.z,1,d))'*Y + lambda*ss.sum_zxy;
s.sum_z = s.z + lambda*ss.sum_z;
s.sum_za = s.z.*s.a + lambda*ss.sum_za;
s.sum_n = 1 + lambda*ss.sum_n;
