% naive bayes, multiple class
function [mean, var, pri] = nb_tr(fea, lab, cn)
% train naive bayes given feature and label

if nargin == 2
    cn = 2;     % 2 classes
elseif nargin < 2
    disp('nb_tr(fea, lab, cn)');
    return;
end

vsize = size(fea,1);
trnum = size(fea, 2);

if cn == 2              % 2 classes
    % calculate prior prob
    pri = zeros(1, 2);
    s1 = sum(lab == 1);
    s2 = sum(lab == -1);
    pri(1) = s1 / trnum;
    pri(2) = s2 / trnum;
    
    % calculate mean
    fea1 = fea(:,find(lab == 1));
    fea2 = fea(:,find(lab == -1));
    
    mean = zeros(vsize, 2);
    mean(:,1) = sum(fea1, 2) / s1;
    mean(:,2) = sum(fea2, 2) / s2;

    % calculate var
    var = zeros(vsize, 2);
    d = fea1 - repmat(mean(:,1), 1, s1);
    var(:,1) = sum(d.*d, 2) / (s1-1);
    d = fea2 - repmat(mean(:,2), 1, s2);
    var(:,2) = sum(d.*d, 2) / (s2-1);
    
else                    % multiple classes
    % calculate prio prob, mean and var
    pri = zeros(1, cn);
    mean = zeros(vsize, cn);
    var = zeros(vsize, cn);
    for k = 1:cn
        sk = sum(lab == k);
        feak = fea(:, find(lab == k));
        pri(k) = sk / trnum;
        mean(:, k) = sum(feak, 2) / sk;
        d = feak - repmat(mean(:,k), 1, sk);
        var(:, k) = sum(d .* d, 2) / (sk-1);
    end
            
end
