function [cluster_idx, cluster_centers] = k_means(input, k, display)
% function [cluster_idx, cluster_centers] = k_means(input, k, display)
% input: NxD matrix, where N is the number of samples and D is the
%        dimension of the samples
% k: number of cluster centers
% display: 0 - do not display figures to show progress
%          1 - will display a figure shows the clustering process, but
%          duplicated colors will appear if k is greater than 7, and only
%          will display when the data is of 2 dimensions or less
% cluster_idx: the corresponding cluster id to each of input
% culster_centers: the final result of k-cluster centers

if ~exist('display', 'var')
    display = 0;
end

[N, D] = size(input);
cluster_idx = zeros(N, 1);

if display && D <= 2
    color = ['y', 'm', 'c', 'r', 'b', 'g', 'k'];
else
    display = 0;
end

% randomly initilize the centers to some datapoints
idx = randperm(N);
cluster_centers = input(idx(1:k), :);
new_cluster_centers = zeros(k, D);

distances = zeros(N, k);
% update centers iteratively
for iter = 1 : 10000
    
    % calculate Euclidean distance of every point to each cluster center
    for i = 1 : k
        diff = bsxfun(@minus, cluster_centers(i,:), input);
        distances(:,i) = sqrt(sum(diff.^2,2));
    end
    
    
    
    % recalculate cluster centers
    [~, cluster_idx] = min(distances, [], 2);
    for i = 1 : k
        new_cluster_centers(i,:) = mean(input(cluster_idx==i,:));
    end
    
    % try to display cluster data using different colors, but no more than
    % 7 colors
    if display
        disp(['iteration', num2str(iter)]);
        figure(1); clf(1);hold on;
        for i = 1 : k
            plot(input(cluster_idx==i,1),input(cluster_idx==i,2),...
                [color(mod(i, length(color))),'.'],'MarkerSize',12);
        end
        plot(cluster_centers(:,1),cluster_centers(:,2),'ko',...
            'MarkerSize',12,'LineWidth',2);
        hold off; 
        drawnow;
        pause(0.5);
    end
    
    % stop updating while centers converged
    if sum(sum((new_cluster_centers - cluster_centers).^2)) < 1e-10
        break;
    end
    cluster_centers = new_cluster_centers;
end
return;

