function [sampleLabels, groupBases, segmentationHistory, mergeLabelHistory] = entropy_gpca(X, epsilon, affine, normalizeDistance)
% Inputs:
%
% X: The set of data vectors to be segmented, arranged in columns.
%
% epsilon: epsilon^2 is the acceptable mean-squared coding error.  If you already have a good estimate of the MSE of your data relative to their subspaces, this is a good value to start.  If you want to scale your data and get the same segmentation, scale epsilon by the same amount.
%
% affine: if true, algorithm will fit a mixture of affine subspaces model, rather than a mixture of linear subspaces model.
%
% normalizeDistance: if true, algorithm will use a steepest descent rule for merges rather than a greedy one.  
%
% Note: This takes only slightly longer, but we think there may be examples where it helps convergence, although we haven't found any yet.  If you find this flag improves segmentation for your application, please tell Drew Wagner (awagner@uiuc.edu) about it.
%
% Outputs:
% sampleLabels: A vector containing the group membership for each sample.
% groupBases: A 

KEEPLOG = true;
if nargout < 3,
    KEEPLOG = false;
    segmentationHistory = [];
    mergeLabelHistory = [];
end

DEBUG = false;  % Only used for debugging purposes.

if nargin < 4,
    normalizeDistance = false;
end

epsilon2=epsilon^2;

[dimensionCount, sampleCount] = size(X);

sampleLabels = 1:sampleCount;
subspaceCount = sampleCount;

% entropySets is a vector whose i-th element is the entropy of the i-th
% group.
squaredNormX = sum(X.^2,1);
if affine
    entropySets = dimensionCount * log2(1 + (squaredNormX/epsilon2))/ 2 + log2(sampleCount);
else
    entropySets = (dimensionCount + 1) * log2(1 + (dimensionCount/ epsilon2)*squaredNormX)/ 2 + log2(sampleCount);
end

% Initialize a table to log the history of the segmentations.  Each row is a segmentation.
if(KEEPLOG),
    segmentationHistory = nan(sampleCount, sampleCount);
    segmentationHistory(1,:) = sampleLabels;
    mergeLabelHistory = nan(sampleCount-1, 2);    
end

if(normalizeDistance),
    % Initialize an array to tabulate slope of objective function for pairwise mergings.
    % This is a square matrix whose ij-th element is slope of the objective function for
    % merging groups i and j into one group.
    % Meaningful values are maintainged on the strict upper triangular portion.  Other portion keept at infinity.
    objectiveArray = inf(subspaceCount);
    for subspaceIndex1 = 1 : subspaceCount-1
        for subspaceIndex2 = subspaceIndex1+1 : subspaceCount
            entropy1 = entropySets(subspaceIndex1);
            entropy2 = entropySets(subspaceIndex2);
            entropyUnion = entropy_function(X(:, [subspaceIndex1 subspaceIndex2]),epsilon2,sampleCount,affine);
            segmentationDistance = vi_incremental(sampleCount, 1, 1);
            objectiveArray(subspaceIndex1, subspaceIndex2) = (entropyUnion - entropy1 - entropy2) / segmentationDistance;
        end;
    end;
else,
    % Initialize an array to tabulate savings in entropy if pairs of groups are merged.
    % objectiveArray is a square matrix whose ij-th element is the entropy lost by
    % merging groups i and j into one group.
    % Meaningful values are maintainged on the strict upper triangular portion.  Other portion keept at infinity.
    objectiveArray = inf(subspaceCount);
    for subspaceIndex1 = 1 : subspaceCount-1
        for subspaceIndex2 = subspaceIndex1+1 : subspaceCount
            entropy1 = entropySets(subspaceIndex1);
            entropy2 = entropySets(subspaceIndex2);
            entropyUnion = entropy_function(X(:, [subspaceIndex1 subspaceIndex2]),epsilon2,sampleCount,affine);
            objectiveArray(subspaceIndex1, subspaceIndex2) = entropyUnion - entropy1 - entropy2;
        end;
    end;
end

% Keep track of the sizes of the groups.
if(normalizeDistance),
    groupSizes = ones(1, sampleCount);  % For a bottom-up segmenation every sample is a group of size one.
end

% Main Loop.  At each turn, apply the best pairwise merge until objective will longer decreases by pairwise merging.
while 1,
    % We are done if there are no more subspaces to be merged together.
    if (subspaceCount==1)
        break;
    end;

    % Find the best two subspaces to merge together.
    % This will be the group that has the best decrease in entropy per amount of segmentation change.
    %   This results in a closer analogy to steepest descent.
    slope = objectiveArray;
    [temp minIndex] = min(objectiveArray, [], 1);
    [minObjective minIndex2] = min(temp);
    minIndex1 = minIndex(minIndex2);

    % If this does not decrease the entropy, we are done.
    if (minObjective > 0)
        break;
    end;

    if(DEBUG > 2 && normalizeDistance),
        groupSizesOld = groupSizes;
        sampleLabelsOld = sampleLabels;
    end
    
    % Merging groups with labels minIndex1 and minIndex2...
    if(DEBUG),
        disp(['Merging groups with indeces ' num2str(minIndex1) ' and ' num2str(minIndex2)]);
    end

    if(KEEPLOG),
        mergeLabelHistory(sampleCount - subspaceCount + 1, :) = [minIndex1 minIndex2];
    end
    
    % Relabel groups so that the second group chosen for merging has the last label, by swapping labels with the last group.
    if (minIndex2 ~= subspaceCount),

        % Update the sample labels vector
        sampleLabels(sampleLabels == subspaceCount) = -1;
        sampleLabels(sampleLabels == minIndex2) = subspaceCount;
        sampleLabels(sampleLabels == -1) = minIndex2;

        % Update the vector storing the entropies of each group.
        [entropySets(minIndex2), entropySets(subspaceCount)] = swap(entropySets(minIndex2), entropySets(subspaceCount));

        % Update the matrix storing the change in objective function (or slope if normalizing) when groups are merged.
        objectiveArray = swap_matrix(objectiveArray, minIndex2, subspaceCount);
        
        % Update the matrix storing the sizes of the groups.
        if(normalizeDistance),
            [groupSizes(minIndex2), groupSizes(subspaceCount)] = swap(groupSizes(minIndex2), groupSizes(subspaceCount));
        end
    end

    % Merging groups with labels minIndex1 and subspaceCount...
    
    % Update vector containing sample labels.
    sampleLabels(sampleLabels == subspaceCount) = minIndex1;

    % Update vector containing group entropies.
    entropySets(subspaceCount) = nan; % Avoids an array resize.
    entropySets(minIndex1) = entropy_function(X(:, sampleLabels == minIndex1),epsilon2,sampleCount,affine);

    % Update the matrix storing the sizes of the groups;
    if(normalizeDistance),
        groupSizes(minIndex1) = groupSizes(minIndex1) + groupSizes(subspaceCount);
        groupSizes(subspaceCount) = nan;  % Avoids an array resize.
    end

    % Update the number of subspaces.
    subspaceCount = subspaceCount - 1;

    % Update the array containing the changes in entropy (or the slope if normalizing).  Note:  BIG, un-necessary array resize.
    % Indexing complicated somewhat by the strict upper-diagonalness of our data structure.
    objectiveArray = objectiveArray(1:end-1, 1:end-1);
    for subspaceIndex2 = [1:minIndex1-1 minIndex1+1:subspaceCount],
        i = min(minIndex1, subspaceIndex2); j = max(minIndex1, subspaceIndex2);
        entropy1 = entropySets(i);
        entropy2 = entropySets(j);
        entropyUnion = entropy_function(X(:, (sampleLabels == i) | (sampleLabels == j)) ,epsilon2,sampleCount,affine);
        
        if(normalizeDistance),
            objectiveArray(i, j) = (entropyUnion - entropy1 - entropy2) / vi_incremental(sampleCount, groupSizes(i), groupSizes(j));
        else,
            objectiveArray(i, j) = entropyUnion - entropy1 - entropy2;
        end
    end;
        
    % Log the current segmentation.
    if(KEEPLOG)
        segmentationHistory(sampleCount - subspaceCount + 1, :) = sampleLabels;
    end
    
end;  % End main loop.

% Use the current segmentation and epsilon to determine the actual subspace models.
[groupMeans groupBases] = subspace_bases(X, sampleLabels, epsilon2, affine);

% Compute the VI between a segmentation and a pairwise merge of the segmentation.
% Needs only the sizes of the groups to be merged and the total number of samples.
function vi = vi_incremental(n, n1, n2);
temp1 = (n1+n2)/n;
temp2 = (n1/n);
temp3 = (n2/n);
vi = temp1*log2(temp1) - temp2*log2(temp2) - temp3*log2(temp3);

% If A is an upper triangular matrix where A(m,n) = f(x_m,x_n) for m < n
% relabel so that the roles of x_i and x_j are swapped.
function A = swap_matrix(A, i, j)
[A(1:i-1, i), A(1:i-1, j)] = swap(A(1:i-1, i), A(1:i-1, j));
[A(i, i+1:j-1), A(i+1:j-1, j)] = swap(A(i, i+1:j-1), A(i+1:j-1, j));
[A(i, j+1:end), A(j, j+1:end)] = swap(A(i, j+1:end), A(j, j+1:end));

% swap the values of x and y
function [y, x] = swap(x,y)