function [Pi,Mu1,Sigma1,Gamk, Gamtheta, L] = EMGauGam(X,ltol,maxiter,pflag,Init, Verbose)
% 
% EM algorithm for the mixture of Gaussian and Gamma estimation
%
% Inputs:
%   X - input data
%   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
%   Mu1 - estimated the mean of Gaussians
%   Sigma1 - estimated the standard deviation the Gaussians
%   Gamk - estimated k parameter for Gamma distribution
%   Gamtheta - estimated theta paprameter for Gamma distribution
%   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.01;
end

if isempty(maxiter)
    maxiter = 500;
end

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

%%%% EM algorithm %%%%
niter = 0;
disp(sprintf('%f %f %f %f %f %f',Mu1,Sigma1,Gamk,Gamtheta));
while niter<=maxiter
    [newL, E] = Expectation(X,Pi,Mu1,Sigma1, Gamk, Gamtheta); % E-step    
    [Pi,Mu1,Sigma1, Gamk, Gamtheta] = Maximization(X,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));
        disp(sprintf('%f %f %f %f %f %f',Mu1,Sigma1,Gamk,Gamtheta));
    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_GM(X,Pi,Mu1,Sigma1,Gamk, Gamtheta);
    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,Pi,Mu1,Sigma1,Gamk, Gamtheta)
E = zeros(size(X,1), 3);
% 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
p1 = Pi(1) * normpdf(X,Mu1,Sigma1);
p2 = Pi(2) * gampdf(X,Gamk,Gamtheta);
E(:,1) = p1./(p1+p2);
E(:,2) = p2./(p1+p2);
%E(:,3) = p3./(p1+p2);

% Compute Likelihood
L = sum(E(:,1) .* ( repmat(log(Pi(1)), size(X)) - repmat(log(2*pi)/2, size(X)) - ...
    repmat(log(Sigma1), size(X)) - (X-Mu1).^2/2*Sigma1^2 ) + ...
    E(:,2) .* (repmat(log(Pi(2)), size(X)) - repmat(Gamk*log(Gamtheta), size(X)) - ...
    repmat(gammaln(Gamk), size(X)) + (Gamk-1) * log(X) - X./Gamtheta));% + repmat(Gamtheta, size(X))) );
%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% End of Expectation %%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function [Pi,Mu1,Sigma1,Gamk, Gamtheta] = Maximization(X,E)
Pi(1) = sum(E(:,1))/size(X,1);
Pi(2) = sum(E(:,2))/size(X,1);

Mu1 = sum(E(:,1).*X)/sum(E(:,1));
Sigma1 = sqrt(sum(E(:,1).*((X-Mu1).^2))/sum(E(:,1)));

temp = log(repmat(sum(E(:,2).*X),size(X))./(X*sum(E(:,2))));
lambda = sum(E(:,2).*temp)/sum(E(:,2));
if lambda >= 0 && lambda <= 0.5772
    Gamk = (0.5000876+0.1648852*lambda-0.0544274*lambda^2)/lambda;
end
if lambda > 0.5772 && lambda <= 17.0
    Gamk = (8.898919+9.059950*lambda-0.9775373*lambda^2)/...
        (lambda*(17.79728+11.968477*lambda+lambda^2));
end
if lambda > 17.0
    Gamk = 1/lambda;
end
if Sigma1 < 0.0001
    Sigma1 = rand();
end

% if Gamk > 10
%     Gamk = rand();
% end
Gamtheta = sum(E(:,2).*X)/ (Gamk*sum(E(:,2)));
if Gamtheta < 0.001;
    Gamtheta = rand();
end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% 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, Mu1, Sigma1, Gamk, Gamtheta] = Init_EM(X)
% [Ci,C] = kmeans(X,3,'Start','cluster', ...
%     '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,3,'Start','cluster', ...
%         'Maxiter',100, ...
%         'EmptyAction','drop', ...
%         'Display','off');
% end
% Pi(1) = length(find(Ci==1))/size(X,1);
% Pi(2) = length(find(Ci==2))/size(X,1);
% Pi(3) = length(find(Ci==3))/size(X,1);
% Mu1 = C(1);
% Mu2 = C(2);
% g1 = X(Ci==1);
% g2 = X(Ci==2);
% Sigma1 = std(g1);
% Sigma2 = std(g2);
% if Sigma1 < 0.0001
%     Sigma1 = rand();
% end
% if Sigma2 < 0.0001
%     Sigma2 = rand();
% end
% g3 = X(Ci==3);
% Gamtheta = var(g3)/mean(g3);
% if Gamtheta < 0.01;
%     Gamtheta = rand();
% end
% Gamk = mean(g3)/Gamtheta;
% if Gamk > 10
%     Gamk = rand();
% end
Pi = zeros(1,2);
[P,M1,S1,M2, S2] = EMGauGau(X,[],500,0,[], false);
M = [M1 M2];
S = [S1 S2];
[Mu2 Mu2idx] = min(M);
Pi(2) = P(Mu2idx);
Sigma2 = S(Mu2idx);
Gamtheta = Sigma2^2/Mu2;
Gamk = Mu2/Gamtheta;

M(Mu2idx) = [];
S(Mu2idx) = [];
P(Mu2idx) = [];
Mu1 = M(1);
Sigma1 = S(1);
Pi(1) = P(1);

if Sigma1 < 0.0001
    Sigma1 = rand();
end
%%%%%%%%%%%%%%%%%%%%%%%%
%%%% End of Init_EM %%%%
%%%%%%%%%%%%%%%%%%%%%%%%

function Plot_GM(X, Pi, Mu1, Sigma1, Gamk, Gamtheta) %W,M,V
[f,r] = ecdf(X);
ecdfhist(f,r)
hold on

Q = zeros(size(r));
P = Pi(1)*normpdf(r,Mu1,Sigma1);
Q = Q + P;
plot(r,P,'r--', 'linewidth',1.5); 
hold on

P = Pi(2)*gampdf(r,Gamk,Gamtheta);
Q = Q + P;
plot(r,P,'y--', 'linewidth',1.5); 
hold on

plot(r,Q,'g-','linewidth',2);
grid on
xlabel('X');
ylabel('Probability density');

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