% This script tests the performance of the basic k-means algorithm on a few
% different test cases. Choose one geometry by setting the "flag_geom" flag
% as follows:
% 
% === flag_geom=1
%     This is a very simple clustering problem with 3 distinct clusters.
%
% === flag_geom=2
%     The points are not distributed on a circle. In this case, the
%     minimization problem has lots of "low-energy" configurations, and
%     takes a while to converge.
%
% === flag_geom=3
%     The points are distributed uniformly in a square. Now, k-means solves
%     a problem of finding (sort of) a Voronoi tessellation problem.
%
% === flag_geom=4
%     This is the "Mickey Mouse" geometry with one large cluster and two
%     small ones. Observe how the minimization problem underlying k-means
%     strongly favors the clusters to be more or less equi-sized. The
%     resulting clustering is not the "natural" one for this example.
%
% /// P.G. Martinsson, April 25, 2016

function main_kmeans02

rng('default')
rng(7)
colvec = 'rgbmcy';

flag_geom = 4;

if (flag_geom == 1)
  nclust = 3;
  xxA = [    0.15*randn(1,400);...
             0.15*randn(1,400)];
  xxB = [1 + 0.15*randn(1,400);...
             0.15*randn(1,400)];
  xxC = [    0.15*randn(1,400);...
         1 + 0.15*randn(1,400)];
  xx = [xxA,xxB,xxC];
  n  = size(xx,2);
  [~,ind] = sort(rand(1,n));
  xx = xx(:,ind);
  %CC = [0.5,-0.05,0.05;...
  %        0,   1,  1];
elseif (flag_geom == 2)
  nclust = 4;
  n      = 1000;
  tt     = 2*pi*rand(1,n);
  xx     = [cos(tt) + 0.2*randn(1,n);...
            sin(tt) + 0.2*randn(1,n)];
elseif (flag_geom == 3)
  nclust = 4;
  n      = 5000;
  xx     = rand(2,n);
elseif (flag_geom == 4)
  nclust = 3;
  n1     = 5000;
  R1     = 0.5;
  n2     = 800;
  R2     = 0.2;
  n3     = 800;
  R3     = 0.2;
  tt1    = 2*pi*rand(1,n1);
  rr1    = R1*sqrt(rand(1,n1));
  tt2    = 2*pi*rand(1,n2);
  rr2    = R2*sqrt(rand(1,n2));
  tt3    = 2*pi*rand(1,n3);
  rr3    = R3*sqrt(rand(1,n3));
  xx     = [rr1.*cos(tt1), 0.45 + rr2.*cos(tt2), -0.45 + rr3.*cos(tt3);...
            rr1.*sin(tt1), 0.55 + rr2.*sin(tt2),  0.55 + rr3.*sin(tt3)];...
  n      = size(xx,2);
  CC     = 0.2*randn(2,nclust);
end
ax = [min(xx(1,:))-0.1,max(xx(1,:))+0.1,min(xx(2,:))-0.1,max(xx(2,:))+0.1];

%%% Initialize centers if this was not done in setting up the geometry
if ~exist('CC','var')
  [~,ind] = sort(rand(1,n));
  CC      = xx(:,ind(1:nclust));
end
J = 1:n;

for istep = 1:100
  fprintf(1,'Iteration %3d\n',istep)
  
  %%% Assignment step
  J_old = J;
  DD1   = ones(nclust,1)*xx(1,:) - CC(1,:)'*ones(1,n);
  DD2   = ones(nclust,1)*xx(2,:) - CC(2,:)'*ones(1,n);
  RRSQ  = DD1.*DD1 + DD2.*DD2;
  [~,J] = min(RRSQ);

  figure(1)
  hold off
  for i = 1:nclust
    Iclust = find(J == i);
    c = colvec(mod(i-1,length(colvec))+1);
    plot(xx(1,Iclust),xx(2,Iclust),['.',c],'MarkerSize',10);
    hold on
    han = plot(CC(1,i),CC(2,i),'ko','MarkerSize',10);
    set(han,'MarkerFaceColor',c,'LineWidth',3)
  end
  axis equal
  axis(ax)
  keyboard

  if (norm(J - J_old) == 0)
    fprintf(1,'Converged!!\n')
    keyboard
  end
  
  %%% Update centers.
  for i = 1:nclust
    Iclust = (J == i);
    CC(:,i) = mean(xx(:,Iclust),2);
  end
  
end

return