function [rbm,output] = trainrbm(input, nh, epochs, e, wc, m)
%Given a series of inputs, trains an RBM with nh hidden units.
%Plotted error is the average across visible units and inputs.
% rbm{1}= weights (w), rbm{2}= hidden biases (hb), rbm{3}= visible biases (vb)
% output= hidden activation probabilities for every input (after training)
% input = [dimesion, batch_size, num_batches]
% nh    = number of hidden units
% e     = learning rate (default 0.01)
% wc    = weight cost   (default 0).   Should be small (e.g. .0002).
% m     = momentum      (default 0).   Range [0,1-e]     (e.g. .9).
if nargin<4; e=0.01; end;  if nargin<5; wc=0; end;  if nargin<6; m=0; end

[nv,batsz,nbat]=size(input);   output = zeros(nh,batsz,nbat);
vb =zeros(nv,1);  hb =zeros(nh,1);  w =.1*randn(nh,nv);
dvb=zeros(nv,1);  dhb=zeros(nh,1);  dw=   zeros(nh,nv);  %d for delta

figure; hold on; xlabel epoch; ylabel 'avg error'; xlim([0 epochs]);
for epoch=1:epochs;
  err=0;  %fprintf(1, 'epoch %3i ', epoch);
  for batch=1:nbat
    %Sample hidden units given input, then reconstruct visible and hidden.
    v = input(:,:,batch);
    h = 1./(1 + exp(-(w *v + repmat(hb,[1 batsz]))));
    hs= h > rand(nh, batsz);
    vr= 1./(1 + exp(-(w'*hs+ repmat(vb,[1 batsz]))));
    hr= 1./(1 + exp(-(w *vr+ repmat(hb,[1 batsz]))));

    %Contrastive Divergence rule: dw ~ h*v - hr*vr
    dw  = m*dw  + e*( (h*v'-hr*vr')/batsz - wc*w);  w = w +dw;
    dvb = m*dvb + (e/batsz)*(sum(v,2)-sum(vr,2));   vb= vb+dvb;
    dhb = m*dhb + (e/batsz)*(sum(h,2)-sum(hr,2));   hb= hb+dhb;
    
    err = err   + sum(sum( (v-vr).^2 ));
    output(:,:,batch)=h;  %
  end
  %fprintf(1, ' error %.4f  \n', err/(nv*batsz*nbat));
  plot( epoch, err/(nv*batsz*nbat), '.'); drawnow;
end
rbm = {w,hb,vb}; %alternatively, rbm = struct('w',w,'hb',hb,'vb',vb);