
function nntrain(epmax, nstat)

% function nntrain(epmax, nstat)
%
% nntrain trains a 3-layer neural net; given X and Y it tries to
% find values for weights W1, W2, W3, and biases B1, B2, and B3,
% such that the expression W3*T(W2*T(W1*X+B1)+B2)+B3 is as close
% to Y as possible.  This is done with a (typically) long series
% of incremental improvements, basically gradient descent with 
% very short steps along the graient.
% 
% parameters
%   epmax  -- limit on training epochs
%   nstat  -- option to print status every nstat epoch
%
% file inputs (set by nninit.m)
%   X, Y         - training matrices 
%   Xmean, Xstd  - X mean and standard deviation
%   Ymean, Ystd  - Y mean and standard deviation
%   W1, W2, W3   - weight matrices 
%   B1, B2, B3   - bias vectors
%
% file outputs (used by nnfwd.m)
%   W1, W2, W3   - weight matrices 
%   B1, B2, B3   - bias vectors
%

% default is not to print a status message
if nargin == 1
  nstat = 0;
end

% backprop training parameters
err_goal        = 0.0001;  % stop if scaled error is <= err_goal
lr              = 0.001;   % initial learning rate
lr_inc          = 1.05;    % learning rate increment
lr_dec          = 0.7;	   % learning rate decrement
mm_const        = 0.95;	   % momentum constant
err_rat         = 1.04;	   % error ratio

% load training data
load X -ascii
load Y -ascii

% load stats on training data
load Xmean -ascii
load Xstd  -ascii
load Ymean -ascii
load Ystd  -ascii

% load the weights and biases
load W1 -ascii
load W2 -ascii
load W3 -ascii
load B1 -ascii
load B2 -ascii
load B3 -ascii

% start timing here
tic

% normalize X
[Xrows,Xcols] = size(X);
for i = 1:Xcols
  X(:,i) = (X(:,i) - Xmean) ./ Xstd;
end
  
% normalize Y
[Yrows,Ycols] = size(Y);
for i = 1:Ycols
  Y(:,i) = (Y(:,i) - Ymean) ./ Ystd;
end

% initial derivatives
dW1 = W1 * 0;
dB1 = B1 * 0;
dW2 = W2 * 0;
dB2 = B2 * 0;
dW3 = W3 * 0;
dB3 = B3 * 0;

% initial momentum 
mm = 0;

% initial forward computation
A1  =  tanh(W1*X + B1*ones(1,Xcols));
A2  =  tanh(W2*A1 + B2*ones(1,Xcols));
A3  =  W3*A2 + B3*ones(1,Xcols);
E   =  Y - A3;
sse =  sum(sum(E .* E));

% backprop to compute deltas
D3 = E ;
D2 = (1 - (A2 .* A2)) .* (W3' * D3) ;
D1 = (1 - (A1 .* A1)) .* (W2' * D2) ;

% loop on epochs
for epoch = 1 : epmax
  
  % scale sse by number of elements
  scaled_err = sse / (Yrows*Ycols);

  % print a periodic error summary
  if mod(epoch, nstat) == 0
    fprintf(1, 'epoch %4d,  scaled error %7.5f,  learning ratio %6.5f\n', ...
                epoch, scaled_err, lr);
  end

  % quit if error is acceptable
  if scaled_err <= err_goal 
    epoch = epoch-1;
    break;
  end

  % compute weight adjustments
  t1 = (1 - mm) * lr;
  dW1 = mm * dW1 + t1 * (D1 * X') ;
  dB1 = mm * dB1 + t1 * (D1 * ones(Xcols,1)) ;

  dW2 = mm * dW2 + t1 * (D2 * A1') ;
  dB2 = mm * dB2 + t1 * (D2 * ones(Xcols,1)) ;

  dW3 = mm * dW3 + t1 * (D3 * A2') ;
  dB3 = mm * dB3 + t1 * (D3 * ones(Xcols,1)) ;

  % try out modified weights
  W1_tmp = W1 + dW1;  B1_tmp = B1 + dB1;
  W2_tmp = W2 + dW2;  B2_tmp = B2 + dB2;
  W3_tmp = W3 + dW3;  B3_tmp = B3 + dB3;

  % forward computation
  A1_tmp  =  tanh(W1_tmp*X + B1_tmp*ones(1,Xcols));
  A2_tmp  =  tanh(W2_tmp*A1_tmp + B2_tmp*ones(1,Xcols));
  A3_tmp  =  W3_tmp*A2_tmp + B3_tmp*ones(1,Xcols);
  E_tmp   =  Y - A3_tmp;
  sse_tmp =  sum(sum(E_tmp .* E_tmp));

  mm = mm_const;

  % compare new and previous error
  if sse_tmp > sse * err_rat

    % new weight and activation set is too much worse
    lr = lr * lr_dec;
    mm = 0;

  else

    % new weight and activation set is acceptable
    if sse_tmp < sse
      % increment learning rate 
      lr = lr * lr_inc;
    end
    
    % save changes
    W1 = W1_tmp;  B1 = B1_tmp;
    W2 = W2_tmp;  B2 = B2_tmp;
    W3 = W3_tmp;  B3 = B3_tmp;
    A1 = A1_tmp;  A2 = A2_tmp;
    E  = E_tmp;  sse = sse_tmp;

    % backprop to compute deltas
    D3 = E ;
    D2 = (1 - (A2 .* A2)) .* (W3' * D3) ;
    D1 = (1 - (A1 .* A1)) .* (W2' * D2) ;
  end

end

% stop timing, print results
fprintf(1, 'nntrain(): time = %.2f\n', toc)

% save the new weights and biases
save W1 W1 -ascii
save W2 W2 -ascii
save W3 W3 -ascii
save B1 B1 -ascii
save B2 B2 -ascii
save B3 B3 -ascii

