% lda training
function [w, w0] = lda_tr(fea, lab, cn)
% train linear hyperplane using feature and lab

if nargin == 2
    cn = 2;
elseif nargin < 2
    disp('lda_tr(fea, lab, [cn])');
    return;
end

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

if cn == 2                  % 2 classes
    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 common var
    d = fea1 - repmat(mean(:,1), 1, s1);
    var1 = diag(d*d');
    d = fea2 - repmat(mean(:,2), 1, s2);
    var2 = diag(d*d');
    sigma = (var1+var2)/(2*trnum-2);

    % calculate w and w0
    w = (mean(:,1) - mean(:,2)) ./ sigma;
    w0 = log(pri(1)/pri(2)) - 0.5 * (mean(:,1)+mean(:,2))' * ((mean(:,1)-mean(:,2)) ./ sigma);

else                        % multiple classes
    
    pri = zeros(1, cn);
    mean = zeros(vsize, cn);
    var = zeros(vsize, cn);
    w = zeros(vsize, cn);
    w0 = zeros(1, 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) = diag(d*d');
    end
    sigma = sum(var, 2) / (k*trnum-k);
    
    for k = 1:cn
        w(:, k) = mean(:,k) ./ sigma;
        w0(k) = -0.5*mean(:,k)'*(mean(:,k)./sigma) + log(pri(k));
    end
end
