%% This function finds the label of the mode of the K nearest neighbors. If 
%% there are multiple modes, then the label with the smallest average 
%% distance is selected. K is the number of neighbors to consider, D is the 
%% distance between a sample and K neighbors, IDX are the indices of the K 
%% neighbors (in regards to the feature vector), and labels_train are the 
%% labels for each sample.

function predicted_labels = findNearestNeighbor(K, D, IDX, labels_train)
    predicted_labels = zeros(size(D));
    for i = 1:length(D)
        neighbor_labels = labels_train(IDX{i}(1:K));
        [~,~, nearestNeighbors] = mode(neighbor_labels);
        % If there is one distinct mode, choose that as the nearest
        % neighbor. Otherwise, choose the most common label with the
        % minimum average distance
        if length(nearestNeighbors{1}) == 1  
            predicted_labels(i) = nearestNeighbors{1};
        else
            avg_dist=zeros(length(nearestNeighbors{1}), 1);
            for j = 1:length(nearestNeighbors{1})
                avg_dist(j) = mean(D{i}(find(neighbor_labels == nearestNeighbors{1}(j))));
            end
            [~,index] = min(avg_dist);
            predicted_labels(i) = nearestNeighbors{1}(index);
        end
    end