% Illustrate use of the MAP-DP algorithm for clustering, Normal-Wishart
% case. Please see "What to do when K-means clustering fails: a simple yet
% principled alternative algorithm" by Yordan P. Raykov et al. for further
% details.

clear all;

% Initialize RNGs for reproduceability
rand('twister',3);
randn('state',24);

D = 2;      % Data dimensionality
K = 3;      % Number of clusters

% Generate random cluster locations and covariance matrices
Sigma = zeros(D,D,K);
Mu = zeros(D,K);
for k = 1:K
    S = randn(D,D);
    S = S'*S;
    Sigma(:,:,k) = 0.15*S;
    Mu(:,k) = 2.8*randn(D,1);
end

% Generate random cluster mixture weights
Pi = rand(1,K);
Pi = Pi/sum(Pi);

N = 1000;   % Number of data points to generate

% Generate categorical data for cluster assignments z
Z = randsample(K,N,true,Pi);

% Generate multivariate Gaussian data for X
X = zeros(D,N);
for k = 1:K
    i = (Z == k);
    M = sum(i);
    X(:,i) = sqrtm(Sigma(:,:,k))*randn(D,M)+repmat(Mu(:,k),1,M);
end

% Randomized restarting (Appendix C)
rand('twister',211631);
i = randperm(N)';
Z = Z(i);
X = X(:,i);

% Set up Normal-Wishart MAP-DP prior parameters
N0 = 0.5;                           % Prior count (concentration parameter)
m0 = mean(X,2);                     % Normal-Wishart prior mean
a0 = 10;                            % Normal-Wishart prior scale
c0 = 10/N;                          % Normal-Wishart prior degrees of freedom
B0 = diag(1./(0.05*var(X,0,2)));    % Normal-Wishart prior precision

% Run Normal-Wishart MAP-DP to convergence
[mu,z,k,E] = S1_Function(X,N0,m0,a0,c0,B0);

% Plot results
close all;
figure;

% Plot generated data
subplot(1,3,1);
hold on;
for j = 1:k
    i = (Z == j);
    col = rand(1,3);
    plot(X(1,i),X(2,i),'.','Color',col);
    plot(Mu(1,j),Mu(2,j),'k+','MarkerSize',20,'LineWidth',2.0);
end
box on;
xlabel('x^1');
ylabel('x^2');
title('Generated data');

% Plot MAP-DP clustering result
subplot(1,3,2);
hold on;
for j = 1:k
    i = (z == j);
    col = rand(1,3);
    plot(X(1,i),X(2,i),'.','Color',col);
    plot(mu(1,j),mu(2,j),'k+','MarkerSize',20,'LineWidth',2.0);
end
box on;
xlabel('x^1');
ylabel('x^2');
title('MAP-DP clustering');

% Plot objective function for each iteration
subplot(1,3,3);
plot(E,'k.-');
box on;
xlabel('Iteration');
ylabel('E');
title('MAP-DP objective function');