function [Pi,Mus,Sigmas, L] = EMKGau(X,K,ltol,maxiter,pflag,Init,Verbose)
% 
% EM algorithm for the mixture of K Gaussians estimation
%
% Inputs:
%   X - input data
%   K - the number of Gaussians
%   ltol - percentage of the log likelihood difference between 2 iterations ([] for none)
%   maxiter - maximum number of iteration allowed ([] for none)
%   pflag - 1 for plotting GM for 1D or 2D cases only, 0 otherwise ([] for none)
%   Init - structure of initial Pi, Mu1, Sigma1, Mu2: Init.Pi, Init.Mu1, Init.Sigma1, Init.Mu2 ([] for none)
%
% Ouputs:
%   Pi - estimated weights of mixture components
%   Mus - estimated means of Gaussians
%   Sigmas - estimated standard deviations the Gaussians
%   L - log likelihood of estimates
%
% April 2009, Keshi Dai
% 
% Credits
% EM_GM - Patrick P. C. Tsui
% http://www.mathworks.com/matlabcentral/fileexchange/8636-emgm
%

if nargin < 6
    Verbose = false;
end

if isempty(ltol)
    ltol = 0.001;
end

if isempty(maxiter)
    maxiter = 500;
end

%%%% Initialize Pi, Mu1, Sigma1, Mu2 %%%%
t = cputime;
if isempty(Init),  
    [Pi, Mus, Sigmas] = Init_EM(X, K);     
else
    Pi = Init.Pi;
    Mus = Init.Mus;
    Sigmas = Init.Sigmas;
    Sigmas(Sigmas < 0.001) = rand();
end
%[newL, E] = Expectation(X,Pi,Mu1,Sigma1,Mu2); % Initialize log likelihood
newL = -Inf;
oldL = 2*newL;

%%%% EM algorithm %%%%
niter = 0;
if K==1
    [Mus, Sigmas] = normfit(X);
    newL = Expectation(X,K,Pi,Mus,Sigmas); % E-step
else
    while niter<=maxiter
        [newL, E] = Expectation(X,K,Pi,Mus,Sigmas); % E-step
        [Pi,Mus,Sigmas] = Maximization(X,K,E);  % M-step
        if abs(newL-oldL)<ltol
            break;
        end
        oldL = newL;
        %newL = Likelihood(X,Pi,Mu1,Sigma1,Mu2);
        niter = niter + 1;
        if Verbose
            disp(sprintf('iterations: %d, log-likelihood: %f',niter-1, newL));
        end
    end
end
L = newL;

%%%% Plot %%%%
if pflag==1,
    if size(X,2)>2
        disp('Can only plot 1 or 2 dimensional applications!\n');
    else
        Plot_KGM(X,K,Pi,Mus,Sigmas);
    end
    elapsed_time = sprintf('CPU time used for EM: %5.2fs',cputime-t);
    if Verbose
        disp(elapsed_time);
        disp(sprintf('Number of iterations: %d',niter-1));
    end
end
%%%%%%%%%%%%%%%%%%%%%%
%%%% End of EM    %%%%
%%%%%%%%%%%%%%%%%%%%%%

function [L, E] = Expectation(X,K,Pi,Mus,Sigmas)
E = zeros(size(X,1), K);
% for n=1:size(X,1)
%     E(n,1) = ComputeGamma(X(n), 1, Pi, Mu1, Sigma1, Mu2);
%     E(n,2) = ComputeGamma(X(n), 2, Pi, Mu1, Sigma1, Mu2);
% end
p = zeros(size(X,1),K);
for i=1:K
    p(:,i) = Pi(i) * normpdf(X,Mus(i),Sigmas(i));
end

for i=1:K
    E(:,i) = p(:,i) ./ sum(p,2);
end

% Compute Likelihood
tempL = zeros(size(X,1),1);
for i=i:K
    tempL = tempL + E(:,i) .* ( repmat(log(Pi(i)), size(X)) - repmat(log(2*pi)/2, size(X)) - ...
    repmat(log(Sigmas(i)), size(X)) - (X-Mus(i)).^2/2*Sigmas(i)^2 );
end
L = sum(tempL);


%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% End of Expectation %%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function [Pi,Mus,Sigmas] = Maximization(X,K,E)
Pi = zeros(1,K);
Mus = zeros(1,K);
Sigmas = zeros(1,K);
for i=1:K
    Pi(i) = sum(E(:,i))/size(X,1);
    Mus(i) = sum(E(:,i).*X)/sum(E(:,i));
    Sigmas(i) = sqrt(sum(E(:,i).*((X-Mus(i)).^2))/sum(E(:,i)));
end
Sigmas(Sigmas < 0.001) = rand();

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% End of Maximization %%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% function L = Likelihood(X,Pi,Mu1,Sigma1,Mu2)
% % Compute L based on K. V. Mardia, "Multivariate Analysis", Academic Press, 1979, PP. 96-97
% % to enchance computational speed
% L = 0;
% for n=1:size(X,1)
%     gamma1 = ComputeGamma(X(n),1,Pi, Mu1, Sigma1, Mu2);
%     gamma2 = ComputeGamma(X(n),2,Pi, Mu1, Sigma1, Mu2);
%     L1 = gamma1 * (log(Pi(1)) - log(2*pi)/2 - log(Sigma1) - (X(n)-Mu1)^2/2*Sigma1^2);
%     L2 = gamma2 * (log(Pi(2)) - X(n)/Mu2 - log(Mu2));
%     L = L + L1 + L2;
% end
%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% End of Likelihood %%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%

% function gamma = ComputeGamma(xn, i, Pi, Mu1, Sigma1, Mu2)
% gamma = 0;
% p1 = Pi(1)*normpdf(xn,Mu1,Sigma1);
% p2 = Pi(2)*exppdf(xn,Mu2);
% if i==1
%     gamma = p1/(p1+p2);
% end
% if i==2
%     gamma = p2/(p1+p2);
% end
%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% End of ComputeGamma %%
%%%%%%%%%%%%%%%%%%%%%%%%%%%

function [Pi, Mus, Sigmas] = Init_EM(X, K)
Pi = zeros(1,K);
Mus = zeros(1,K);
Sigmas = zeros(1,K);
% if (length(X)>50)

[Ci,C] = kmeans(X,K,'Start','sample', ...
    'Maxiter',100, ...
    'EmptyAction','drop', ...
    'Display','off'); % Ci(nx1) - cluster indeices; C(k,d) - cluster centroid (i.e. mean)
while sum(isnan(C))>0,
    [Ci,C] = kmeans(X,K,'Start','sample', ...
        'Maxiter',100, ...
        'EmptyAction','drop', ...
        'Display','off');
end
% else
% 	Ci = zeros(size(X));
% 	half = floor(length(Ci)/2);
% 	Ci(1:half) = 1;
% 	Ci(half:end) = 2;
% 	C = [mean(Ci(1:half)), mean(Ci(half,end))];
% end
for i=1:K
    Pi(i) = length(find(Ci==i))/size(X,1);
    Mus(i) = C(i);
    g = X(Ci==i);
    Sigmas(i) = std(g);
end
Sigmas(Sigmas < 0.001) = rand();
%%%%%%%%%%%%%%%%%%%%%%%%
%%%% End of Init_EM %%%%
%%%%%%%%%%%%%%%%%%%%%%%%

function Plot_KGM(X,K,Pi,Mus,Sigmas) %W,M,V
[f,r] = ecdf(X);
ecdfhist(f,r)
hold on

Q = zeros(size(r));

for i=1:K
    P = Pi(i)*normpdf(r,Mus(i),Sigmas(i));
    Q = Q + P;
    plot(r,P,'r--', 'linewidth',1.5);
    hold on
end

plot(r,Q,'g-','linewidth',2);
grid on
xlabel('X');
ylabel('Probability density');
xlim([min(r), max(r)]);

title('Mixture Models estimated by EM');
%%%%%%%%%%%%%%%%%%%%%%%%
%%%% End of Plot_GM %%%%
%%%%%%%%%%%%%%%%%%%%%%%%