% This code implements the diverse setwise low margin active selection method for
% selecting image samples used to train ranking functions. If you would like to
% use the code for publication, please cite the following paper:
%
%    Lucy Liang and Kristen Grauman. "Beyond Comparing Image Pairs: Setwise Active
%    Learning for Relative Attributes." In CVPR, 2014.

function selectedIndices = DSLM(features, w, sampleSize, k, clusterLabels)

% This function takes, among other things, a set of feature vectors and the current ranking
% function, and returns a set of samples selected from the features for annotation.
% Input:
%    features - n x p matrix representing n feature vectors in p-dimensional feature space.
%    w - current ranking function.
%    sampleSize - the number of samples to return for annotation.
%    k - number of total clusters.
%    clusterLabels - nx1 vector representing the cluster membership for each feature vector in features.
%
% Output:
%    selectedIndices - a sampleSize x 1 vector of indices corresponding to the features selected
%        for annotation.

% First, apply current w (ranking function) to the features.
% featureRankLabels is a nx1 vector of rank labels corresponding to each feature.
featureRankLabels = features * w;

featuresCount = size(features, 1);

% Now actually sort the feature vectors.
sortedIndices = (1:featuresCount)';
featureRankLabelsTemp = [featureRankLabels, sortedIndices];
featureRankLabelsTemp = sortrows(featureRankLabelsTemp, 1);
featureRankLabels = featureRankLabelsTemp(:, 1);  % The computed rankings, sorted.
sortedIndices = featureRankLabelsTemp(:, 2)'; % The indices of the feature vectors in features - in sorted order.
clear featureRankLabelsTemp;

% Generate clusters if necessary.
if clusterLabels == 0
    fprintf('Running kmeans...\n');
    tic;
    clusterLabels = kmeans(features, k);
    toc;
end

% Sort cluster labels as well.
clusterLabels = clusterLabels(sortedIndices);

% Compute every pairwise rank margins - stored in a nxn matrix.
featureRankLabelsCols = repmat(featureRankLabels, 1, length(featureRankLabels));
featureRankLabelsRows = repmat(featureRankLabels', length(featureRankLabels), 1);
allPairwiseRankMargins = abs(featureRankLabelsCols-featureRankLabelsRows);
clear featureRankLabelsCols featureRankLabelsRows;

% Now we can start our iterative process of selecting the best samples!

%Initialize C_prev_base to the first sampleSize number of cumulative rank margins.
tempMat0 = allPairwiseRankMargins(1:sampleSize, 1:sampleSize);
C_prev_base = sum(sum(tempMat0))/2;
clear tempMat0;

C_min = 9999;                       % The current minimum cumulative rank margin.
I_min = (1:sampleSize);             % The indices corresponding to the feature vectors
                                    %     that currently make up the minimum cumulative
                                    %     rank margin.
j_max = featuresCount - sampleSize;

fprintf('Running diverse setwise low margin...\n');
tic;
for j=1:j_max
    j_vec = (j:j+sampleSize-1); % j_vec - vector of sampleSize indices corresponding to
                                %         the current set of feature vectors we're evaluating.
    C_prev = C_prev_base;
    
    while 1
        if C_prev < C_min && C_prev ~= 0 ...
                && length(unique(featureRankLabels(j_vec))) == sampleSize
            % Found new lowest cumulative rank margin. Check that all samples are
            % from different clusters.
            if length(unique(clusterLabels(j_vec))) == sampleSize
		        % Each sample in this set are from different clusters, update
		        % current minimum rank margin info.
                C_min = C_prev;
                I_min = j_vec;
                break;
            else
                % Samples are not all from different clusters - there are
                % duplicate cluster memberships somewhere.

                % Get the next index over. We will replace one of the dupes with
                % this new member.
                nextIndex = max(j_vec) + 1;
                if nextIndex > featuresCount
                    break;
                end

		        % Now find the dupes.
                u = unique(clusterLabels(j_vec));
                h = histc(clusterLabels(j_vec), u);
                dupClusters = u(h > 1);
                dupCluster = dupClusters(1); % Deal with one set of dupes at a time. For
                                             % sampleSize = 4, there can be no more than
                                             % 2 sets of dupes per sample set.
                dupClusterIndices = find(clusterLabels(j_vec) == dupCluster);
                
                % Get rid of one of the dupes.
                % We'll choose the one that, when replaced by the new member,
                % will yield the lowest cumulative rank margin.
 		        dupCount = length(dupClusterIndices);
                costVector = zeros(1, dupCount);
                for t=1:dupCount
                    j_vec_less_t = j_vec;
                    j_vec_less_t(dupClusterIndices(t)) = [];
                    
                    % C_dj is the cost (in rank margin distance) of this current dupe
                    % (indexed at dupClusterIndices(t))
                    C_dj = sum(allPairwiseRankMargins(j_vec(dupClusterIndices(t)), j_vec_less_t));
                    
                    % C_dj_plus is the cost of the new item at next index with
                    % respect to the remaining sampleSize-1 members if the
                    % current dupe is removed.
                    C_dj_plus = sum(allPairwiseRankMargins(nextIndex, j_vec_less_t));
                    
                    % costVector(t) is the new cumulative rank margin if the
                    % current dupe is removed. Lower is better.
                    costVector(t) = C_prev - C_dj + C_dj_plus;
                end
                
                C_prev = min(costVector);

                % Now replace this dupe with the index of the new member.
                j_vec(dupClusterIndices(costVector == C_prev)) = nextIndex;
                
            end
        else
            break;
        end
    end

    % If we broke out the while loop, then either we:
    %    (1) Finished obtaining a sample set with the new lowest
    %        cumulative rank margin, starting an index j, whose
    %        members belong to different clusters or
    %    (2) Weren't able to find any sample set, starting at index j,
    %        whose members belonged to different clusters.
    % In either case, we're done with possible sample sets starting
    % at index j, so update C with the next contiguous data members,
    % starting at j+1.
    C_dj = sum(allPairwiseRankMargins(j, j+1:j+sampleSize-1));
    C_dj_plus = sum(allPairwiseRankMargins(j+sampleSize, j+1:j+sampleSize-1));
    C_prev_base = C_prev_base - C_dj + C_dj_plus;
end
toc;

% Retrieve the indices of the samples to be returned for annotation.
selectedIndices = (sortedIndices(I_min));