function res = Boost(mode, learn, varargin)
% usage is one of the two following:
%   boost = Boost('train', @learner, X, Y, numRounds, learner-args)
%   Y     = Boost('predict', @learner, boost, X, numRounds) -- numRounds is optional
  
  switch lower(mode),
    case 'name',
      res = 'Boost';
      
    case 'train',
      if nargin < 5,
        error('Boost training requires at least five arguments');
      end;
      
      [NX,DX] = size(varargin{1});
      [NY,DY] = size(varargin{2});
      
      if NX ~= NY,
        error('Boost: there must be the same number of data points in X and Y');
      end;
      if DY ~= 1,
        error('Boost: Y must have only one column');
      end
      if ~all(varargin{2}==0 | varargin{2}==1),
        error('boosting only defined for binary (0/1) problems');
      end;

      res = BoostTrain(learn, varargin{1}, varargin{2}, varargin{3}, varargin{4:end});
      
      
    case 'predict',
      % check arguments
      if nargin ~= 4 && nargin ~= 5,
        error('Boost: prediction requires three (or four) arguments: the learner, model, and X (and maybe numRounds)');
      end;
      boost = varargin{1};
      X     = varargin{2};
      maxRounds = length(boost.alphas);
      if nargin == 5 && varargin{3} < maxRounds, 
        maxRounds = varargin{3}; 
      end;
    
      % compute predictions
      res = BoostPredict(learn, boost, X, maxRounds);
  
    otherwise
      error('unknown mode');
  end;
      
  
function boost = BoostTrain(learn, X, Y, numRounds, varargin)
  boost = {};
  boost.models = {};
  boost.alphas = {};
  boost.errors = {};
  
  N = length(Y);
  
  W = ones(N,1) / N;
  
  for t=1:numRounds,
    fprintf('+'); myflush;
    %Wei: Train weak learner using distribution D_t.
    model = learn('train', X, Y, W, 1);
    
    % compute hypothesis
    h = learn('predict', model, X);
    
    % compute error
    e = sum(W(Y ~= h));
    % Need double check. My error decreased with t. But Hal's code
    % increased with t. but my code predicted correctly. Why?
  
    % store error
    boost.errors{t} = e;
    
    % if error is too big, get out
    if e >= 0.5,
      break;
    end;
    
    % compute alpha
    a = 0.5 * log((1-e)/e);
    
    % update weights
    % W = W .* exp((-a) .* Y .* h);   
    for i = 1:N
        W(i) = W(i) * exp(-a * Y(i) * h(i));
    end;
    
    
    % Need double check if there is a '2' at Hal's notes.
    W = W/sum(W);
    
    % store model and alpha
    boost.models{t} = model;
    % By default this is h. Just changed it. To store only h makes no sense
    % because that's for the training data. In prediction we need 'model'
    % to predict for each h_t. Need to make sure this understanding is
    % correct. (Hal wrote strange code here??)
    boost.alphas{t} = a;
  end;

function Y = BoostPredict(learn, boost, X, maxRounds)
  N = size(X,1);
  Y = zeros(N,1);

  for t=1:maxRounds,
    % compute predictions using boost.models{t}:
    h = learn('predict', boost.models{t}, X);
    P = h * boost.alphas{t};
        
    % update overall prediction (note: you may need to ensure P is 0/1 or
    % +1/-1 as you prefer)
    Y = Y + P;
  end;
  % convert to binary output
  Y = Y > 0.5;
  
  
  