% [Projection, Summary] = jPCA(Data, analyzeTimes, params)
%
% OUTPUTS:
%   Projection is a struct with one element per condition.
%   It contains the following fields:
%       .proj           The projection into the 6D jPCA space.
%       .times          Those times that were used
%       .projAllTimes   Pojection of all the data into the jPCs (which were derived from a subset of the data)
%       .allTimes       All the times (exactly the same as Data(1).times)
%       .tradPCAproj    The traditional PCA projection (same times as for 'proj')
%       .tradPCAprojAllTimes   Above but for all times.
%
%   Summary contains the following fields:
%       .jPCs           The jPCs in terms of the PCs (not the full-D space)
%       .PCs            The PCs (each column is a PC, each row a neuron)
%       .jPCs_highD     The jPCs, but in the original high-D space.  This is just PCs * jPCs, and is thus of size neurons x numPCs
%       .varCaptEachJPC The data variance captured by each jPC
%       .varCaptEachPC  The data variance captured by each PC
%       .R2_Mskew_2D    Fit quality (fitting dx with x) provided by Mskew, in the first 2 jPCs.
%       .R2_Mbest_2D    Same for the best M
%       .R2_Mskew_kD    Fit quality (fitting dx with x) provided by Mskew, in all the jPCs (the number of which is set by 'numPCs'.
%       .R2_Mbest_kD    Same for the best M
%
%       There are some other useful currently undocumented (but often self-explanatory) fields in
%       'Summary'
%
%   Summary also contains the following, useful for projecting new data into the jPCA space.  Also
%   useful for getting from what is in 'Projection' back into the high-D land of PSTHs
%   Note that during preprocessing we FIRST normalize, and then (when doing PCA) subtract the
%   overall mean firing rate (no relationship to the cross-condition mean subtraction).  For new
%   data you must do this in the same order (using the ORIGINAL normFactors and mean firing rats.
%   To get from low-D to high D you must do just the reverse: add back the mean FRs and the multiply
%   by the normFactors.
%       .preprocessing.normFactors = normFactors;  
%       .preprocessing.meanFReachNeuron = meanFReachNeuron; 
%
% INPUTS:
% The input 'Data' needs to be a struct, with one entry per condition.
% For a given condition, Data(c).A should hold the data (e.g. firing rates).
% Each column of A corresponds to a neuron, and each row to a timepoint.
%
% Data(c).times is an optional field.  If you provide it, only those entries that match
% 'analyzeTimes' will be used for the analysis. If  analyzeTimes == [], all times will be used.
%  
%  If you don't provide it, a '.times' field
% is created that starts at 1.  'analyzeTimes' then refers to those times.
% 
% 'params' is optional, and can contain the following fields:
%   .numPCs        Default is 6. The number of traditional PCs to use (all jPCs live within this space)
%   .normalize     Default is 'true'.  Whether or not to normalize each neurons response by its FR range.
%   .softenNorm    Default is 10.  Determines how much we undernormalize for low FR neurons.  0 means
%                  complete normalization.  10 means a neuron with FR range of 10 gets mapped to a range of 0.5.
%   .meanSubtract  Default is true.  Whether or not we remove the across-condition mean from each
%                  neurons rate.
%   .suppressBWrosettes    if present and true, the black & white rosettes are not plotted
%   .suppressHistograms    if present and true, the blue histograms are not plotted
%   .suppressText          if present and true, no text is output to the command window 
%
%  As a note on projecting more data into the same space.  This can be done with the function
%  'projectNewData'.  However, if you are going to go this route then you should turn OFF
%  'meanSubtract'.  If you wish to still mean subtract the data you should do it by hand yourself.
%  That way you can make a principled decision regarding how to treat the original data versus the
%  new-to-be-projected data.  For example, you may subtract off the mean manually across 108
%  conditions.  You might then leave out one condition and compute the jCPA plane (with meanSubtract set to false).  
%  You could then project the remaining condition into that plane using 'projectNewData'.
% 
% 
function [Projection, Summary] = jPCA(Data, analyzeTimes, params)


%% quick bookkeeping

numConds = length(Data);
numTimes = size(Data(1).A,1);
% there is also a 'numAnalyzedTimes' defined below.

%% setting parameters that may or may not have been specified

% 'times' field
% if the user didn't specify a times field, create one that starts with '1'
if ~isfield(Data(1),'times')
    for c = 1:length(Data)
        Data(c).times = 1:numTimes;
    end
end

if exist('analyzeTimes', 'var') && ~isempty(analyzeTimes) && max(diff(analyzeTimes)) > max(diff(Data(1).times))
    disp('error, you can use a subset of times but you may not skip times within that subset');
    Projection = []; Summary = []; return;
end

% the number of PCs to look within
numPCs = 6;
if exist('params', 'var') && isfield(params,'numPCs')
    numPCs = params.numPCs;
end
if rem(numPCs,2)>0
    disp('you MUST ask for an even number of PCs.'); return;
end

% do we normalize
normalize = true;
if exist('params', 'var') && isfield(params,'normalize')
    normalize = params.normalize;
end

% do we soften the normalization (so weak signals stay smallish)
% numbers larger than zero mean soften the norm.
% The default (10) means that 10 spikes a second gets mapped to 0.5, infinity to 1, and zero to zero.
% Beware if you are using data that isn't in terms of spikes/s, as 10 may be a terrible default
softenNorm = 10;
if exist('params', 'var') && isfield(params,'softenNorm')
    softenNorm = params.softenNorm;
end

% do we mean subtract
meanSubtract = true;
if exist('params', 'var') && isfield(params,'meanSubtract')
    meanSubtract = params.meanSubtract;
end
if length(Data)==1, meanSubtract = false; end  % cant mean subtract if there is only one condition

if ~exist('analyzeTimes', 'var') || isempty(analyzeTimes)
    disp('analyzing all times');
    analyzeTimes = Data(1).times;
end

%% figure out which times to analyze and make masks
%
analyzeIndices = ismember(round(Data(1).times), analyzeTimes);
if size(analyzeIndices,1) == 1
    analyzeIndices = analyzeIndices';  % orientation matters for the repmat below
end
analyzeMask = repmat(analyzeIndices,numConds,1);  % used to mask bigA
if diff( Data(1).times(analyzeIndices) ) <= 5
    disp('mild warning!!!!: you are using a short time base which might make the computation of the derivative a bit less reliable');
end

% these are used to take the derivative
bunchOtruth = true(sum(analyzeIndices)-1,1);
maskT1 = repmat( [bunchOtruth;false],numConds,1);  % skip the last time for each condition
maskT2 = repmat( [false;bunchOtruth],numConds,1);  % skip the first time for each condition

if sum(analyzeIndices) < 5
    disp('warning, analyzing few or no times');
    disp('if this wasnt your intent, check to be sure that you are asking for times that really exist');
end

%% make a version of A that has all the data from all the conditions.
% in doing so, mean subtract and normalize

bigA = vertcat(Data.A);  % append conditions vertically

% note that normalization is done based on ALL the supplied data, not just what will be analyzed
if normalize  % normalize (incompletely unless asked otherwise)
    ranges = range(bigA);  % For each neuron, the firing rate range across all conditions and times.
    normFactors = (ranges+softenNorm);
    bigA = bsxfun(@times, bigA, 1./normFactors);  % normalize
else
    normFactors = ones(1,size(bigA,2));
end
 


sumA = 0;
for c = 1:numConds
    sumA = sumA + bsxfun(@times, Data(c).A, 1./normFactors);  % using the same normalization as above
end
meanA = sumA/numConds;
if meanSubtract  % subtract off the across-condition mean from each neurons response
    bigA = bigA-repmat(meanA,numConds,1);
end



%% now do traditional PCA

smallA = bigA(analyzeMask,:);
[PCvectors,rawScores] = pca(smallA);  % apply PCA to the analyzed times
meanFReachNeuron = mean(smallA);  % this will be kept for use by future attempts to project onto the PCs

% these are the directions in the high-D space (the PCs themselves)
if numPCs > size(PCvectors,2)
    disp('You asked for more PCs than there are dimensions of data');
    disp('Giving up');
end
PCvectors = PCvectors(:,1:numPCs);  % cut down to the right number of PCs

% CRITICAL STEP
% This is what we are really after: the projection of the data onto the PCs
Ared = rawScores(:,1:numPCs);  % cut down to the right number of PCs

% Some extra steps
%
% projection of all the data
% princomp subtracts off means automatically so we need to do that too when projecting all the data
bigAred = bsxfun(@minus, bigA, mean(smallA)) * PCvectors(:,1:numPCs); % projection of all the data (not just teh analyzed chunk) into the low-D space.
% need to subtract off the mean for smallA, as this was done automatically when computing the PC
% scores from smallA, and we want that projection to match this one.

% projection of the mean

meanAred = bsxfun(@minus, meanA, mean(smallA)) * PCvectors(:,1:numPCs);  % projection of the across-cond mean (which we subtracted out) into the low-D space.

% will need this later for some indexing
numAnalyzedTimes = size(Ared,1)/numConds;

%% GET M & Mskew
% compute dState, and use that to find the best M and Mskew that predict dState from the state

% we are interested in the eqn that explains the derivative as a function of the state: dState/dt = M*State
dState = Ared(maskT2,:) - Ared(maskT1,:);  % the masks just give us earlier and later times within each condition
preState = Ared(maskT1,:);  % just for convenience, keep the earlier time in its own variable

% first compute the best M (of any type)
% note, we have converted dState and Ared to have time running horizontally
M = (dState'/preState');  % M takes the state and provides a fit to dState
% Note on sizes of matrices:
% dState' and preState' have time running horizontally and state dimension running vertically 
% We are thus solving for dx = Mx.
% M is a matrix that takes a column state vector and gives the derivative

% now compute Mskew using John's method
% Mskew expects time to run vertically, transpose result so Mskew in the same format as M
% (that is, Mskew will transform a column state vector into dx)
Mskew = skewSymRegress(dState,preState)';  % this is the best Mskew for the same equation

%% USE Mskew to get the jPCs

% get the eigenvalues and eigenvectors
[V,D] = eig(Mskew); % V are the eigenvectors, D contains the eigenvalues
evals = diag(D); % eigenvalues

% the eigenvalues are usually in order, but not always.  We want the biggest
[~,sortIndices] = sort(abs(evals),1,'descend');
evals = evals(sortIndices);  % reorder the eigenvalues
evals = imag(evals);  % get rid of any tiny real part
V = V(:,sortIndices);  % reorder the eigenvectors (base on eigenvalue size)

% Eigenvalues will be displayed to confirm that everything is working
% unless we are asked not to output text
if ~exist('params', 'var') || ~isfield(params,'suppressText') || ~params.suppressText
    disp('eigenvalues of Mskew: ');
    for i = 1:length(evals)
        if evals(i) > 0;
            fprintf('                  %1.3fi', evals(i));
        else
            fprintf('     %1.3fi \n', evals(i));
        end
    end
end

jPCs = zeros(size(V));
for pair = 1:numPCs/2
    vi1 = 1+2*(pair-1);
    vi2 = 2*pair;
    
    VconjPair = V(:,[vi1,vi2]);  % a conjugate pair of eigenvectors
    evConjPair = evals([vi1,vi2]); % and their eigenvalues
    VconjPair = getRealVs(VconjPair,evConjPair);
    
    jPCs(:,[vi1,vi2]) = VconjPair;
end

%% Get the projections

proj = Ared * jPCs;
projAllTimes = bigAred * jPCs;
tradPCA_AllTimes = bsxfun(@minus, bigA, mean(smallA)) * PCvectors;  % mean center in exactly the same way as for the shorter time period.
crossCondMeanAllTimes = meanAred * jPCs;

% Do some annoying output formatting.
% Put things back so we have one entry per condition
index1 = 1;
index2 = 1;
for c = 1:numConds
    index1b = index1 + numAnalyzedTimes -1;  % we will go from index1 to this point
    index2b = index2 + numTimes -1;  % we will go from index2 to this point
    
    Projection(c).proj = proj(index1:index1b,:);
    Projection(c).times = Data(1).times(analyzeIndices);
    Projection(c).projAllTimes = projAllTimes(index2:index2b,:);
    Projection(c).allTimes = Data(1).times;
    Projection(c).tradPCAproj = Ared(index1:index1b,:);
    Projection(c).tradPCAprojAllTimes = tradPCA_AllTimes(index2:index2b,:);
    
    index1 = index1+numAnalyzedTimes;
    index2 = index2+numTimes;
end
   
%% Done computing the projections, plot the rosette

% do this unless params contains a field 'suppressBWrosettes' that is true
if ~exist('params', 'var') || ~isfield(params,'suppressBWrosettes') || ~params.suppressBWrosettes
    plotRosette(Projection, 1);  % primary plane
    plotRosette(Projection, 2);  % secondary plane
end


%% SUMMARY STATS
%% compute R2 for the fit provided by M and Mskew

% R2 Full-D
fitErrorM = dState'- M*preState';
fitErrorMskew = dState'- Mskew*preState';
varDState = sum(dState(:).^2);  % original data variance

R2_Mbest_kD = (varDState - sum(fitErrorM(:).^2)) / varDState;  % how much is explained by the overall fit via M
R2_Mskew_kD = (varDState - sum(fitErrorMskew(:).^2)) / varDState;  % how much by is explained via Mskew

% unless asked to not output text
if ~exist('params', 'var') || ~isfield(params,'suppressText') || ~params.suppressText
    fprintf('%% R^2 for Mbest (all %d dims):   %1.2f\n', numPCs, R2_Mbest_kD);
    fprintf('%% R^2 for Mskew (all %d dims):   %1.2f  <<---------------\n', numPCs, R2_Mskew_kD);
end


% R2 2-D primary jPCA plane
fitErrorM_2D = jPCs(:,1:2)' * fitErrorM;  % error projected into the primary plane
fitErrorMskew_2D = jPCs(:,1:2)' * fitErrorMskew;  % error projected into the primary plane
dState_2D = jPCs(:,1:2)' * dState'; % project dState into the primary plane
varDState_2D = sum(dState_2D(:).^2); % and get its variance

R2_Mbest_2D = (varDState_2D - sum(fitErrorM_2D(:).^2)) / varDState_2D;  % how much is explained by the overall fit via M
R2_Mskew_2D = (varDState_2D - sum(fitErrorMskew_2D(:).^2)) / varDState_2D;  % how much by is explained via Mskew

if ~exist('params', 'var') || ~isfield(params,'suppressText') || ~params.suppressText
    fprintf('%% R^2 for Mbest (primary 2D plane):   %1.2f\n', R2_Mbest_2D);
    fprintf('%% R^2 for Mskew (primary 2D plane):   %1.2f  <<---------------\n', R2_Mskew_2D);
end

%% variance catpured by the jPCs
origVar = sum(sum( bsxfun(@minus, smallA, mean(smallA)).^2));
varCaptEachPC = sum(Ared.^2) / origVar;  % this equals latent(1:numPCs) / sum(latent)
varCaptEachJPC = sum((Ared*jPCs).^2) / origVar;
varCaptEachPlane = reshape(varCaptEachJPC, 2, numPCs/2);
varCaptEachPlane = sum(varCaptEachPlane);


%% Analysis of whether things really look like rotations (makes plots)

for jPCplane = 1:2
    phaseData = getPhase(Projection, jPCplane);  % does the key analysis
    
    if exist('params', 'var')
        cstats = plotPhaseDiff(phaseData, params, jPCplane);  % plots the histogram.  'params' is just what the user passed, so plots can be suppressed
    else
        cstats = plotPhaseDiff(phaseData, [], jPCplane);
    end
    
    if jPCplane == 1
        circStats = cstats;  % keep only for the primary plane
    end
end

%% Make the summary output structure

Summary.jPCs = jPCs;
Summary.PCs = PCvectors;
Summary.jPCs_highD = PCvectors * jPCs;
Summary.varCaptEachJPC = varCaptEachJPC;
Summary.varCaptEachPC = varCaptEachPC;
Summary.varCaptEachPlane = varCaptEachPlane;
Summary.Mbest = M;
Summary.Mskew = Mskew;
Summary.fitErrorM = fitErrorM;
Summary.fitErrorMskew = fitErrorMskew;
Summary.R2_Mskew_2D = R2_Mskew_2D;
Summary.R2_Mbest_2D = R2_Mbest_2D;
Summary.R2_Mskew_kD = R2_Mskew_kD;
Summary.R2_Mbest_kD = R2_Mbest_kD;
Summary.circStats = circStats;
Summary.acrossCondMeanRemoved = meanSubtract;
Summary.crossCondMean = crossCondMeanAllTimes(analyzeIndices,:);
Summary.crossCondMeanAllTimes = crossCondMeanAllTimes;
Summary.preprocessing.normFactors = normFactors;  % Used for projecting new data from the same neurons into the jPC space
Summary.preprocessing.meanFReachNeuron = meanFReachNeuron; % You should first normalize and then mean subtract using this (the original) mean
% conversely, to come back out, you must add the mean back on and then MULTIPLY by the normFactors
% to undo the normalization.

%% DONE (only functions below)


%% Inline function that gets the real analogue of the eigenvectors
function Vr = getRealVs(V,evals)

    % get real vectors made from the eigenvectors
    
    % by paying attention to this order, things will always rotate CCW
    if abs(evals(1))>0  % if the eigenvalue with negative imaginary component comes first
        Vr = [V(:,1) + V(:,2), (V(:,1) - V(:,2))*1i]; 
    else
        Vr = [V(:,2) + V(:,1), (V(:,2) - V(:,1))*1i];
    end
    Vr = Vr / sqrt(2);

    % now get axes aligned so that plan is spread mostly along the horizontal axis
    testProj = (Vr'*Ared(1:numAnalyzedTimes:end,:)')'; % just picks out the plan times
    rotV = pca(testProj);
    crossProd = cross([rotV(:,1);0], [rotV(:,2);0]);
    if crossProd(3) < 0, rotV(:,2) = -rotV(:,2); end   % make sure the second vector is 90 degrees clockwise from the first
    Vr = Vr*rotV; 

    % flip both axes if necessary so that the maximum move excursion is in the positive direction
    testProj = (Vr'*Ared')';  % all the times
    if max(abs(testProj(:,2))) > max(testProj(:,2))  % 2nd column is the putative 'muscle potent' direction.
        Vr = -Vr;
    end
end

end
%% end of main function


%% Some inline functions are below

%% Plot the rosette itself (may need to spruce this up and move to a subfunction)
function plotRosette(Proj, whichPair)

    d1 = 1 + 2*(whichPair-1);
    d2 = d1+1;

    numConds = length(Proj);

    figure;

    % first deal with the ellipse for the plan variance (we want this under the rest of the data)
    planData = zeros(numConds,2);
    for c = 1:numConds
        planData(c,:) = Proj(c).proj(1,[d1,d2]);
    end
    planVars = var(planData);
    circle([0 0], 2*planVars.^0.5, 0.6*[1 1 1], 1); hold on;
    %fprintf('ratio of plan variances = %1.3f (hor var / vert var)\n', planVars(1)/planVars(2));

    allD = vertcat(Proj(:).proj);  % just for getting axes
    allD = allD(:,d1:d2);
    mxVal = max(abs(allD(:)));
    axLim = mxVal*1.05*[-1 1 -1 1];
    arrowSize = 5;
    for c = 1:numConds
        plot(Proj(c).proj(:,d1), Proj(c).proj(:,d2), 'k');
        plot(Proj(c).proj(1,d1), Proj(c).proj(1,d2), 'ko', 'markerFaceColor', [0.7 0.9 0.9]);

        penultimatePoint = [Proj(c).proj(end-1,d1), Proj(c).proj(end-1,d2)];
        lastPoint = [Proj(c).proj(end,d1), Proj(c).proj(end,d2)];
        arrowMMC(penultimatePoint, lastPoint, [], arrowSize, axLim);

    end

    axis(axLim);
    axis square;
    plot(0,0,'k+');
    
    title(sprintf('jPCA plane %d', whichPair));
end


%% Getting the phases
function phaseData = getPhase(Proj, whichPair)
    numConds = length(Proj);
    d1 = 1 + 2*(whichPair-1);
    d2 = d1+1;
    
    for c=1:numConds
        data = Proj(c).proj(:,[d1,d2]);
        phase = atan2(data(:,2), data(:,1));  % Y comes first for atan2
        
        deltaData = diff(data);
        phaseOfDelta = atan2(deltaData(:,2), deltaData(:,1));  % Y comes first for atan2
        phaseOfDelta = [phaseOfDelta(1); phaseOfDelta];  %#ok<AGROW> % so same length as phase
        radius = sum(data.^2,2).^0.5;
        
        % collect and format
        % make things run horizontally so they can be easily concatenated.
        phaseData(c).phase = phase'; %#ok<AGROW>
        phaseData(c).phaseOfDelta = phaseOfDelta'; %#ok<AGROW>
        phaseData(c).radius = radius'; %#ok<AGROW>
        
        % angle between state vector and Dstate vector
        % between -pi and pi
        phaseData(c).phaseDiff = minusPi2Pi(phaseData(c).phaseOfDelta - phaseData(c).phase); %#ok<AGROW>
    end
    
end



%% plotting the phase difference between dx(t)/dt and x(t) where x is the 2D state
function circStatsSummary = plotPhaseDiff(phaseData, params, jPCplane)
    % compute the circular mean of the data, weighted by the r's
    circMn = circ_mean([phaseData.phaseDiff]', [phaseData.radius]');
    resultantVect = circ_r([phaseData.phaseDiff]', [phaseData.radius]');
    
    
    bins = pi*(-1:0.1:1);
    cnts = histc([phaseData.phaseDiff], bins);  % not for plotting, but for passing back out
    

    % do this unless params contains a field 'suppressHistograms' that is true
    if ~exist('params', 'var') || ~isfield(params,'suppressHistograms') || ~params.suppressHistograms
        figure;
        hist([phaseData.phaseDiff], bins, 'Normalization', 'pdf'); hold on;
        plot(circMn, 1, 'ro', 'markerFa', 'r', 'markerSiz', 8);
        plot(pi/2*[-1 1], [0 0], 'ko', 'markerFa', 'r', 'markerSiz', 8);
        set(gca,'XLim',pi*[-1 1]);
        title(sprintf('jPCs plane %d', jPCplane));
    end
    
    %fprintf('(pi/2 is %1.2f) The circular mean (weighted) is %1.2f\n', pi/2, circMn);
    
    % compute the average dot product of each datum (the angle difference for one time and condition)
    % with pi/2.  Will be one for perfect rotations, and zero for random data or expansions /
    % contractions.
    avgDP = averageDotProduct([phaseData.phaseDiff]', pi/2);
    %fprintf('the average dot product with pi/2 is %1.4f  <<---------------\n', avgDP);
    
    circStatsSummary.circMn = circMn;
    circStatsSummary.resultantVect = resultantVect;
    circStatsSummary.avgDPwithPiOver2 = avgDP;  % note this basically cant be <0 and definitely cant be >1
    circStatsSummary.DISTRIBUTION.bins = bins;
    circStatsSummary.DISTRIBUTION.binCenters = (bins(1:end-1) + bins(2:end))/2;
    circStatsSummary.DISTRIBUTION.counts = cnts(1:end-1);
    circStatsSummary.RAW.rawData = [phaseData.phaseDiff]';
    circStatsSummary.RAW.rawRadii = [phaseData.radius]';
    
end

%% for computing the average dot product with a comparison angle
function avgDP = averageDotProduct(angles, compAngle, varargin)
    
    x = cos(angles-compAngle);
    
    if ~isempty(varargin)
        avgDP = mean(x.*varargin{1}) / mean(varargin{1});  % weighted sum
    else
        avgDP = mean(x);
    end
end