function [theta,py] = nbtrain(Xtr,ytr,ess,p)
% train regularized naive bayes parameters
% ------------------------------------------
% X,y: data to train
% ess: equivalent sample size
% p:   prior probability

[r c] = size(Xtr);

if nargin<4  p = mean(ytr);      end;
if nargin<3  ess = floor(r/10);  end;

N(2) = sum(ytr);  % number of docs labelled 1
N(1) = r - N(2);  % number of docs labelled 0

py(1) = N(1)/r;
py(2) = N(2)/r;

theta = zeros(2,c);

ynz = find(ytr);
yze = find(ytr==0);

for i=1:c
    theta(1,i) = sum(Xtr(yze,i));
    theta(2,i) = sum(Xtr(ynz,i));
    
    theta(1,i) = (theta(1,i)+ess*(1-p)) / (N(1)+ess);
    theta(2,i) = (theta(2,i)+ess*p) / (N(2)+ess);
end;
