function [C,fval,exitflag,output] = train(C,Y,X,P,R,S,F,I,T,options)
% TRAIN Learn the parameters using MAP on training data
%
% [C fval exitflag output] = TRAIN(C,Y,X,P,R,S,F,I,T,options) 
%
% Input parameters:
% [Required]
%
%  C:  Input SDMaxEnt classifier object. Training starts from the
%      weights of C
%  Y:  A cell of images (matrices) containing a training labels in
%      1:C.numLabels at every pixel.
%  X:  A cell of feature images (3-D matrices) with the first two
%      dimensions the same size as the corresponding entry in Y and
%      feature vectors along the third dimension.
%
% [Optional]
%
%  P:  Prior on weights for MAP-based training, which may be one of:
%     - 'none':         No prior, uses Maximum Likelihood training.
%     - 'gauss':        IID Gaussian (L2) [Default]. 
%     - 'laplace':      IID Laplacian (L1). Performs feature
%                       selection according to Williams [2].
%  R:  Regularization values (prior parameters).
%      For 'gauss' it is the standard deviation (default 10) and
%      'laplace' it is the scale parameter alpha (default 0.1).
%  S:  Feature selection. A struct of logical (binary) matrices the
%      same size as C.weights (use function featureSelection)
%      indicating the weights that should be learned. Use an empty
%      matrix [] as a placeholder to use all weights but specify further
%      parameters.
%  F:  Indicates whether the selected features are frozen at their
%      current values (F==1) or are given zero values (F==0).
%  I:  Cell containing weights of the training pixels in the
%      log-likelihood objective. Size is the same as for Y. Use an empty
%      matrix [] as a placeholder to equally weight instances but
%      specify further parameters. 
%  T:  Indicates whether the parameterization is complete (T=1) or
%      overcomplete (T=0) [default]. If complete, the system will
%      be faster since one less set of convolutions need to be computed.
%
%  options: Optimization struct passed directly to the
%           optimizer. See function LBFGS.
%
% Outputs:
%
% C:         The learned classifier object (the same, save for updated
%            weights)
%  fval:     Value of the objective function upon start and termination.
%  exitflag: Exit value of the optimizer. See function lbfgs.
%  output:   Output struct from the optimizer. See function lbfgs.
%
%
% References:
% 
% [1] W. Buntine and A. Weigend. Bayesian back-propagation. Complex
%     Systems, 5:603{643, 1991.
%
% [2] P. M. Williams. Bayesian regularization and pruning using a
%     Laplace prior. Neural Computation, 7:117{143, 1995.

  
% SDMAXENT.TRAIN
	
% $Id: train.m 1 2009-05-07 13:41:00Z weinman $
 
% Jerod Weinman
% jerod@acm.org
% Copyright 2007, 2008, 2011
 
%    This file is part of Matlab SDMaxEnt.
%
%    Matlab SDMaxEnt is free software: you can redistribute it and/or modify
%    it under the terms of the GNU General Public License as published by
%    the Free Software Foundation, either version 3 of the License, or
%    (at your option) any later version.
%
%    Matlab SDMaxEnt is distributed in the hope that it will be useful,
%    but WITHOUT ANY WARRANTY; without even the implied warranty of
%    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
%    GNU General Public License for more details.
%
%    You should have received a copy of the GNU General Public License
%    along with Matlab SDMaxEnt.  If not, see <http://www.gnu.org/licenses/>.
  
% Handle arguments and set defaults
  if (nargin>=4)
	
	if strcmpi(P,'gauss')
	  P = 2;
	elseif strcmpi(P,'laplace');
	  P = 1;
	elseif strcmpi(P,'none');
	  P = 0;
	else
	  error(sprintf('Unrecognized regularization method ''%s''',P));
	end;
  else
	P = 0;
  end;
  
  if (nargin<5)
	switch P
	 case 0
	  R = 0;
	 case 1
	  R = .1;
	 case 2
	  R = 10;
	end;
  else
	if R<0
	  error('Negative regularization parameter is not permitted');
	end;
	
	
  end;
  
  
  if (nargin<6)
	S = [];% No feature selection specified. Use all features/weights.
  end;

  if nargin<7
	F = 0;% Default behavior: Zero weights unselected.
  end;
  
  if nargin<8 | isempty(I)
	I = []; % No instance weighting specified. Equally weight instances.
  elseif length(I)~= length(X)
	error('Instance weights length do not match number of instances');
  elseif any(cellfun(@(w) any(w(:)<0),I))
	error('Instance weights must be non-negative');
  end;
  
  if nargin<9
    T = 0; % Overcomplete default.
  end;
  
  if nargin<10
	options = [];% Use the default optimization parameters
  end;
	
  if length(Y) ~= length(X)
	error('Labels and feature graphs lengths do not match');
  end;

  % Enforce the feature selection "zeroing" behavior
  if ~isempty(S) & ~F
	for f=fieldnames(S)'
	  C.weights.(f{:}) = C.weights.(f{:}) .* S.(f{:});
	end;
  end;

  
  % Pass in training data and params to the value function
  value(C.params,Y,X,P,R,S,F,I,T);

  % Get hook to the value function
  F = @value;

  % Set persistent variables in parameter vector converter
  vector2params(C.params);
  W = params2vector(C.params);

  
  
  tic;
  % Get the initial value
  [V,G] = value(W);
  toc;
  
  if isempty(options)
    options = struct('Display','iter',...
                     'MaxIter',1000,...
                     'TolFun',1e-4,...
                     'TolX', 1e-10,...
                     'TolGrad',1e-4,...
                     'Convex',0);
  end;		
  
  % Optimize starting at initial weights Wv and the function F = @value
  % L-BFGS is an approximate second-order convex optimization method.
  
  [W,fval,exitflag,output] = lbfgs( F, W, options);
	
  
  

  % Set the weights we just received from the optimizer  
  C.params = vector2params(W);
 
  
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function W = params2vector(P)
  
  W = [P.weights(:)' P.bias(:)'];
  
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function P = vector2params(W)
  
  persistent M K L
  
  if isa(W,'struct')
	szW = size(W.weights);
	K = szW(1:2);
	M = szW(3);
	L = szW(4);
	return;
  end;
  
  p0 = prod([K M L]);


  P.weights = reshape(W(1:p0),[K M L]);
  P.bias = reshape(W(p0+1:end),[1 1 L]);
  
  
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [V,G] = value(varargin)
  
  persistent Y X P R S F I T Wold Vold Gold ;
  
  switch length(varargin) 
   case 1
	W = varargin{1}; % weights VECTOR
   case 9
	
	% Set the values of the persistent variables
	[W,Y,X,P,R,S,F,I,T] = varargin{:};
	
	% params struct
	% labels
	% data
	% regularization prior
	% regularization prior params
	% feature selection
	% freeze unselected weights
	% instance weights
    % compleTe==1 or overcompleTe==0 (default)
    
	Wold = [];
	% And exit
	return;
	
   otherwise
	error('Incorrect arguments');
  end;
    
  if ~isempty(Wold) & max(abs(Wold-W))<5*eps
	V = Vold;
	G = Gold;
	return;
  else
	Wold = W; % cache the vector
	W = vector2params(W); % Get the param struct
  end;
  
  % Feature selection
  if ~isempty(S) & ~F
	for f=fieldnames(S)'
	  W.(f{:}) = W.(f{:}) .* S.(f{:});
	end;
  end;
  
  szW = size(W.weights);

  K = szW(1:2); % kernel dimensions
  H = szW(3);   % number of channels
  L = szW(4);   % number of labels
	
  V = 0;
  
  sz0 = size(Y{1});

  Gw = zeros(size(W.weights),class(W.weights)); % kernel/weight gradient
  Gb = zeros(size(W.bias),class(W.bias));       % bias gradient
  
  for a=1:length(Y)

	xx = X{a};
	
	sz = size(Y{a});
	
	B = zeros([sz(1:2)-szW(1:2)+1 L]);
	
	for c=1:L
      
      if c==L && T
        break;
      end;
      
	  uu = conv2(xx(:,:,1),W.weights(:,:,1,c),'valid');
	  
	  for h=2:H
		uu = uu + conv2(xx(:,:,h),W.weights(:,:,h,c),'valid');
	  end;

	  B(:,:,c) = uu+W.bias(c);
	  
	end;
	
	mx = max(B,[],3);
	lgZ = log(sum(exp(B-mx(:,:,ones(1,L))),3)) + mx;
	
	B = B - lgZ(:,:,ones(1,L));
	
	Yc = kcrop(Y{a},K);

	[rr,cc] = meshgrid(1:size(B,1),1:size(B,2));	  
	ndx = sub2ind(size(B),rr(:),cc(:),Yc(:));

    if ~isempty(I) && ~isempty(I{a})
      Ic = kcrop(I{a},K);
      ndxI = sub2ind(size(Ic),rr(:),cc(:));
      
      V = V + sum(Ic(ndxI).*B(ndx));
    else
      Ic = [];
      V = V + sum(B(ndx));
    end;


	E = -exp(B);
	E(ndx) = E(ndx) + 1;
	
	[gw,gb] = calcGradient(W,E,xx,Ic,T);
	
	Gw = Gw + gw;
	Gb = Gb + gb;
	
	sz0 = sz;
	
  end;
  
  G.weights = krot(Gw);
  G.bias = Gb;
  
  
  [rV,rG] = regValue(G,W,P,R);

  V = V+rV;
  
  for f=fieldnames(G)'
	G.(f{:}) = G.(f{:}) + rG.(f{:});
	
	if ~isempty(S)
	  G.(f{:}) = G.(f{:}) .* S.(f{:});
	end;
	
  end;
  
  G = params2vector(G);
  
  % Finally, since the opt toolbox does minimization
  V = -V;
  G = -G;

  
  Vold = V;
  Gold = G;

  
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 
function [Gw,Gb] = calcGradient(W,E,X,I,T)
  

  Gw = W.weights*0;
  Gb = W.bias*0;

  szW = size(Gw);

  L = szW(4);
  
  sz = size(X);  

  
  for c=1:L

    if c==L && T
      break;
    end;
    if isempty(I)    
      Er = kadd(rot90(E(:,:,c),2),sz(1:2));
    else
      Er = kadd(rot90(I.*E(:,:,c),2),sz(1:2));
    end;
	
	for h=1:szW(3)
	  Gw(:,:,h,c) = conv2(Er,X(:,:,h),'valid');
	end;
	

  end;

  if isempty(I)
    Gb = sum(sum(E,1),2);
  else
    Gb = sum(sum(I(:,:,ones(1,L)).*E,1),2);
  end;


  if T
    Gb(end) = 0;
  end;
  

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function C = kcrop(Y,K)
  
  hs = K/2;
  B = [ceil(hs)-1 ; floor(hs)];

  
  C = Y(B(1,1)+1:end-B(2,1),B(1,2)+1:end-B(2,2));
  
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function C = kadd(A,D)
  
  szA = size(A);
  
  dSz = D - szA;
  
  C = zeros(szA+2*dSz);
  
  
  C(dSz(1)+1:end-dSz(1),dSz(2)+1:end-dSz(2)) = A;
  
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function K = krot(K)
  
  for h=1:size(K,3)
	for k=1:size(K,4)
	  K(:,:,h,k) = rot90(K(:,:,h,k),2);
	end;
  end;
  
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [V,G] = regValue(G,W,P,R)
% REGVALUE Regularization value and gradient
  
  
  rG = G;
  
  for f=fieldnames(G)'
	% Regularization
	switch P
	 case 0
	  rv = 0;
	  rG.(f{:})(:) = 0;
	 case 1
	  rv = - sumv(abs(W.(f{:})))*R;
	  rG.(f{:}) = - sign(W.(f{:}))*R;
	 case 2
	  rv = - sumv(W.(f{:}).^2)/(2*R^2);
	  rG.(f{:}) = - W.(f{:})/(R^2);
	 otherwise
	  error('Unknown regularization');
	end;
  end;
  
  V = rv;
  G = rG;
