function [center, U, distortion, allCenter] = vqKmeans(data, clusterNum, plotOpt)
% vqKmeans: Vector quantization using K-means clustering (Forgy's batch-mode method)
%	Usage: [center, U, distortion, allCenter] = vqKmeans(data, clusterNum, plotOpt)
%		data (dim x dataNum): data set to be clustered; where each column is a sample data
%		clusterNum: number of clusters (greater than one), or matrix of columns of centers 
%		center (dim x clusterNum): final cluster centers, where each column is a center
%		U: final partition matrix
%		distortion: values of the objective function during iterations 
%
%	Roger Jang, 20030330

if nargin<1, selfdemo; return; end
if nargin<3, plotOpt=0; end

dim=size(data,1);
dataNum=size(data,2);
if clusterNum==1 | (size(clusterNum,1)>1 & size(clusterNum,2)==1)
	center=mean(data, 2);
	U=ones(1, dataNum);
	distortion=sum(sum((data-center*ones(1,dataNum)).^2));
	allCenter=center;
	return;
end

maxLoopCount = 100;				% Max. iteration
distortion = zeros(maxLoopCount, 1);		% Array for objective function
if length(clusterNum)==1
	center = vqInitCenter(data, clusterNum, 3);	% Initial cluster centers
else
	center = clusterNum;				% The passed argument is actually a matrix of cluster centers
end
if nargout==4, allCenter(:,:,1)=center; end

if plotOpt & dim==2
	plot(data(1,:), data(2,:), 'b.');
	centerH=line(center(1,:), center(2,:), 'color', 'r', 'marker', 'o', 'linestyle', 'none', 'linewidth', 2);
	axis image
end;

% Main loop
for i = 1:maxLoopCount,
	[center, distortion(i), U] = vqUpdateCenter(center, data);
	if nargout==4, allCenter(:,:,i+1)=center; end
	if plotOpt, fprintf('Iteration count = %d/%d, distortion = %f\n', i, maxLoopCount, distortion(i)); end
	if plotOpt & dim==2 
		set(centerH, 'xdata', center(1,:), 'ydata', center(2,:));
		drawnow;
	end
	% check termination condition
	if (i>1) & (abs(distortion(i-1)-distortion(i))<eps), break; end
end
loopCount = i;	% Actual number of iterations 
distortion(loopCount+1:maxLoopCount) = [];
if plotOpt & dim==2, figure; vqDataPlot(data, center); end

% ========== subfunctions ==========

% ====== Update centers
function [center, distortion, U] = vqUpdateCenter(center, data)
dim = size(data, 1);
dataNum = size(data, 2);
centerNum = size(center, 2);
% ====== Compute distance matrix
distMat=pairwiseSqrDistance(center, data);
% ====== Find the U (partition matrix)
[minDist, colIndex] = min(distMat);
U = zeros(size(distMat));
U(colIndex+centerNum*(0:dataNum-1)) = 1;
distortion = sum(minDist);	% objective function
% ====== Find new centers
index=find(sum(U,2)==0);
emptyGroupNum=length(index);
if emptyGroupNum==0	% Find the new centers (with no empty cluster)
	center = (data*U')./(ones(dim,1)*sum(U,2)');	% Find the new centers
else	% Add new centers for the empty clusters
	fprintf('Found %d empty group(s)!\n', emptyGroupNum);
	U(index,:)=[];
	center = (data*U')./(ones(dim,1)*sum(U,2)');	% Find the new centers
	if emptyGroupNum<=centerNum/2	% т