function [dprime ROC w criterion classAcc] = caoLDA(cluster1,cluster2,graphics,nTrain);
% Usage:  [dprime ROC w] = caoLDA(cluster1,cluster2,[graphics],[nTrain]);
%
% Calculates Fischer linear discriminant, described by w which is a vector
% of coefficients in n-dimensional space, where the number of dimensions is
% determined by the number of columns in cluster1 & cluster2 
%
% If no arguments are passed, caoLDA makes up some clusters in 3D and shows
% the clouds of points, the discriminant axis, the projection of those
% points to the axis, calculated d' and the ROC curve.
%
% Optional graphics flag turns off display and line reports.
%
% More details for LDA (aka cannonical variate analysis)
%
% The Fischer Linear discriminant finds the line in n-dimensional space
% onto which two clusters of data can be projected in a way that maximizes
% the between-class (cluster) variance with respect to the within class 
% variance.

% The vector w represents the coefficients of a linear function in n
% dimenstions onto which we project the data. The projected data from each
% cluster have a mean m = (1/n)SUMi(yi) where yi = w'xi (the magnitude of
% the projection of each data point, xi, in a cluster onto the unit vector
% w). The scatter of the projected yi's is si^2 = SUMi(yi-m)^2. The Fischer
% Linear Discriminant is the linear funcsion w that maximizes 
% J(w) = |m1^2-m2^2|^2 / (s1^2 + s2^2). (Each of those m's and s's should
% have a tilda over it to indicate that they are means and scatters of the
% projected data points.) Duda, Hart and Stork (2nd edition) derive in 
% Ch. 3.8.2 that the solution is:  w = Sw-1(m1 - m2), where now Sw, m1 and
% m2 are measured on the original data.  Sw is the the sum of S1 and S2, the
% within-class scatter matrices for the data.
%
% cao 11/2009 ... updated 5/2011 to add criterion & leave-n-out classification

if nargin < 3; graphics = 1; end

if nargin < 1
    % Start by making 2 clusters in 3D.
    mu1 = [3 4 3]; sig1 = [1 1 1];
    mu2 = [4 4 3.5]; sig2 = [1 1 1];
    nSamples = 30;
    for n = 1:3
        cluster1(:,n) = repmat(mu1(n),[nSamples 1]) + sig1(n)*randn([nSamples 1]);
        cluster2(:,n) = repmat(mu2(n),[nSamples 1]) + sig2(n)*randn([nSamples 1]);
    end
    subplot(2,1,1)
    plot3(cluster1(:,1),cluster1(:,2),cluster1(:,3),'k.'); hold on
    plot3(cluster2(:,1),cluster2(:,2),cluster2(:,3),'r.'); hold off
else
    nSamples1 = size(cluster1,1);
    nSamples2 = size(cluster2,1);
    if nSamples1 ~= nSamples2
        error('caoLDA written to assume same number of samples in each cluster.')
    else
        nSamples = nSamples1;
    end
    nDims = size(cluster1,2);
    if graphics
        disp(['Number of dimensions:  ' num2str(nDims) '.'])
    end
    if size(cluster2,2) ~= nDims
        error('Input clusters have different number of dimensions (columns).')
    end
    if nDims > nSamples-1
        disp('nDims > nSamples; using PCA to reduce dimensionality')
        [U S V] = svd(cluster1); 
        cluster1 = U*S(:,1:(nSamples-1))*V(1:(nSamples-1),1:(nSamples-1))';
        [U S V] = svd(cluster2); 
        cluster2 = U*S(:,1:(nSamples-1))*V(1:(nSamples-1),1:(nSamples-1))';
    end
end
    
% Duda, Hart and Stork (2nd edition) derive in
% Ch. 3.8.2 that the solution is:  w = Sw-1(m1 - m2), where now Sw, m1 and
% m2 are measured on the original data.  Sw is the the sum of S1 and S2, the
% within-class scatter matrices for the data:
m1 = mean(cluster1); m2 = mean(cluster2);
S1 = (cluster1 - repmat(m1,[nSamples 1]))'*(cluster1 - repmat(m1,[nSamples 1]));
S2 = (cluster2 - repmat(m2,[nSamples 1]))'*(cluster2 - repmat(m2,[nSamples 1]));

% now w is the n-dimensional normal to the discriminant hyperplane
w = -nSamples*inv(S1 + S2)*(m1 - m2)';
w = w/sqrt(sum(w.^2)); % make it a unit vector for fun

if nargin < 1
    plotLim = [min([cluster1;cluster2]);max([cluster1;cluster2])];
    hold on; plot3(w(1)*plotLim(:,1),w(2)*plotLim(:,2),w(3)*plotLim(:,3),'b'); hold off
    axis equal
end

%% Value of discriminant can be estimated from D' for the projected data
projection1 = cluster1*w;
projection2 = cluster2*w;
range = [min(min(projection1,projection2)) max(max(projection1,projection2))];
bins = range(1):(range(2)-range(1))/50:range(2);
p1 = hist(projection1,bins); 
p2 = hist(projection2,bins); 

stdev1 = std(projection1);
stdev2 = std(projection2);
diff = abs(mean(projection2) - mean(projection1));
dprime = 2*diff/(stdev1 + stdev2);
if graphics
    if nargin < 1; subplot(2,2,3); else subplot(1,2,1); end
    bar(bins,p1,'k'); hold on; bar(bins,p2,'r'); hold off
    title(['d'' = ' num2str(dprime,2)])
    xlabel('Projection to discriminant axis')
end

%% Now make an ROC curve to learn a thing or two
for iB = 1:length(bins)
    % classify each point as being from group 2 if greater than criterion
    hits = projection2 > bins(iB);
    false_alarms = projection1 > bins(iB);
    ROC(iB,:) = [length(find(false_alarms)) length(find(hits))];
end
ROC = ROC/nSamples;
dist_from_ideal = sqrt((ROC(:,1) - 0).^2 + (ROC(:,2) - 1).^2);
[temp idxC] = min(dist_from_ideal);
criterion = bins(idxC);
if graphics
    if nargin < 1; subplot(2,2,3); else; subplot(1,2,1); end
    hold on; plot([criterion criterion],[0 max(p1)],'g','linewidth',2); hold off
    if nargin < 1; subplot(2,2,4); else; subplot(1,2,2); end
    plot(ROC(:,1),ROC(:,2))
    xlabel('false alarms'); ylabel('hits')
end

%% Now do what we do w/ fMRI data ... leave-one-out train & classify
if exist('nTrain','var')
    correct1 = []; correct2 = []; z1 = []; z2 = [];
    if graphics; figure; end
    for iT = 1:ceil(100/(nSamples-nTrain))
        randIdx1 = randperm(size(cluster1,1));
        randIdx2 = randperm(size(cluster2,1));
        % sample subset of the data:
        cluster1t = cluster1(randIdx1(1:nTrain),:);
        cluster2t = cluster2(randIdx2(1:nTrain),:);
        m1 = mean(cluster1t); m2 = mean(cluster2t);
        w = -(nSamples-1)*inv(cov(cluster1t) + cov(cluster2t))*(m1 - m2)';
        w = w/sqrt(sum(w.^2)); 
        % use ROC curve to develop criterion that maximizes hits while minimizing false alarms
        % (a more traditional criterion would be developed as the point
        % equidistant between the 2 cluster means ... comment out the next
        % 10 lines and uncomment 157 if you want the other criterion
        for iB = 1:length(bins)
            % classify each point as being from group 2 if greater than criterion
            hits = cluster2t*w > bins(iB);
            false_alarms = cluster1t*w > bins(iB);
            ROC(iB,:) = [length(find(false_alarms)) length(find(hits))];
        end
        ROC = ROC/(nSamples-1);
        dist_from_ideal = sqrt((ROC(:,1) - 0).^2 + (ROC(:,2) - 1).^2);
        [temp idxC] = min(dist_from_ideal);
        criterion = bins(idxC);
       % criterion = (mean(cluster1t*w) + mean(cluster2t*w))/2;
        correct1 = [correct1 ; cluster1(randIdx1((nTrain+1):end),:)*w < criterion];
        correct2 = [correct2 ; cluster2(randIdx2((nTrain+1):end),:)*w > criterion];
        if graphics
            subplot(2,1,1)
            projection1 = cluster1t*w;
            projection2 = cluster2t*w;
            range = [min(min(projection1,projection2)) max(max(projection1,projection2))];
            bins = range(1):(range(2)-range(1))/50:range(2);
            p1 = hist(projection1,bins); 
            p2 = hist(projection2,bins); 
            bar(bins,p1,'k'); hold on; bar(bins,p2,'r'); hold off
            hold on; plot([criterion criterion],[0 max(p1)],'g','linewidth',2')
            plot(cluster1(randIdx1((nTrain+1):end),:)*w,zeros([1 length((nTrain+1):nSamples)]),'k*')
            plot(cluster2(randIdx2((nTrain+1):end),:)*w,zeros([1 length((nTrain+1):nSamples)]),'r*');hold off
            z1 = [z1; (cluster1(randIdx1((nTrain+1):end),:)*w - mean(cluster1t*w))/std(cluster1t*w)];
            z2 = [z2; (cluster2(randIdx2((nTrain+1):end),:)*w - mean(cluster2t*w))/std(cluster2t*w)];
        end
    end
    if graphics
        subplot(2,2,3); hist(z1,20); title('z-scores, category 1')
        subplot(2,2,4); hist(z2,20); title('z-scores, category 2')
    end
    disp(['Classification for category 1: ' num2str(100*sum(correct1)/length(correct1)) '% correct'])
    disp(['Classification for category 2: ' num2str(100*sum(correct2)/length(correct2)) '% correct'])
    classAcc = [sum(correct1)/length(correct1) sum(correct2)/length(correct2)];
else
    classAcc = NaN;
end