%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%% Make a Linear Discriminant Classifier 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function [fractionCorrect,classifierOrClasses] = LDA(data,CV,selectInOneTrialF,command)
%%% [fractionCorrect,classifierOrClasses] = LDA(train,CV)
%%%   train.features: ntasks x ntrials cell array of featurized data
%%%        .classes: ntasks x ntrials matrix of class labels
%%%   CV contains for collectTrainValidation()
%%%     .parms.firstMode
%%%     .parms.nModes
%%%     .validateIndices: ntasks x ntrials cell array of indices
%%%     .classes:  ntasks x ntrials cell array of class labels (ints)
%%%   command:  Optional. If present, must be 'make' and classifier is created.
%%% Returns
%%%   classifier.weights  
%%%             .bias
%%%   validateCorrect: fraction of validation samples correct
%%%   X: training data

if nargin > 3 && strcmp(command,'make')

  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  %%% make the classifier
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

  if ~isfield(CV.parms,'firstMode')
    CV.parms.firstMode = 1;  %so this works if data.features are not cell array of
    CV.parms.nModes = 1;    %windows, but just nchannels x nsamples matrix
  end

  [X,Y,Xval,Yval] = collectPartitions(data,CV,selectInOneTrialF,'train');

  [p,N] = size(X);
  K = size(data.features,1);

  %% project to first fisherDim Fisher Dimensions
  %% normalize data
  X = X';
  c = cov(X);
  sq = inv(sqrtm(c));
  means = mean(X);
  Xs = (X - repmat(means,size(X,1),1)) * sq;
  Yindicator = zeros(length(Y),K);
  for i=1:length(Y)
    Yindicator(i,Y(i)) = 1;
  end
  A = ccakirby(Xs,Yindicator(:,1:end-1));
  X = X * A;
  X = X(:,1:CV.parms.fisherDim);
  X = X';
  classifier.means = means;
  classifier.sq = sq;
  classifier.proj = A;
  classifier.fisherDim = CV.parms.fisherDim;
  
  Xval = Xval';
  Xval = (Xval - repmat(means,size(Xval,1),1)) * sq * A;
  Xval = Xval(:,1:CV.parms.fisherDim);
  Xval = Xval';
  
  p = CV.parms.fisherDim;
  
  %% Calculate class priors, means and sum of covariance matrices.
  priors = zeros(1,K);
  means = zeros(K,p);
  covsum = zeros(p,p);
  classNames = unique(Y);
  for i = 1:length(classNames)
    mask = Y == classNames(i);
    Nthisclass = sum(mask);
    priors(i) = Nthisclass / N;
    Xmat = X(:,mask);
    means(i,:) = mean(Xmat');
    XmeanZero = Xmat - repmat(means(i,:)',1,Nthisclass);
    covsum = covsum + XmeanZero * XmeanZero';
  end
  covariance = covsum / (N-K);
  invcov = inv(covariance);

  %% Calculate weights and biases for each class
  weights = zeros(p,K);
  bias = zeros(K,1);
  for k = 1:K
    weights(:,k) = invcov * means(k,:)';
    bias(k) = - 0.5 * means(k,:) * weights(:,k) + log(priors(k));
  end

  classifier.weights = weights;
  classifier.bias = bias;
  
  if ~isempty(CV.validateIndices)
    disc = classifier.weights' * Xval + repmat(classifier.bias,1,size(Xval,2));
    [junk,whichmax] = max(disc);
    classes = classNames(whichmax);
    validateCorrect = sum(classes == Yval) / length(Yval);
  else
    validateCorrect = [];
  end

  classifierOrClasses = classifier;
  fractionCorrect = validateCorrect;

elseif nargin > 3 && ~strcmp(command,'make')

  error('Command to LDA must be ''make'' or absent.');

else
  
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  %%% use the classifier
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

  [X,Y] = collectPartitions(data,CV,selectInOneTrialF);

  %% fisher
  X = X';
  X = (X - repmat(CV.classifier.means,size(X,1),1)) * CV.classifier.sq * CV.classifier.A;
  X = X(:,1:CV.classifier.fisherDim);
  X = X';

  disc = CV.classifier.weights' * X + repmat(CV.classifier.bias,1,size(X,2));
  [junk,whichmax] = max(disc);
  classNames = unique(Y);
  classes = classNames(whichmax);
  testCorrect = sum(classes == Y) / length(Y);

  classifierOrClasses = classes;
  fractionCorrect = testCorrect;

end
