function em_demo

% data
a = 2;
b = -1.5;

% parameter range
x = -3:.02:4;

% figure properties
figure(1)
set(gca, 'FontSize', 15)

% likelihood function
plot(x, lik(x,a,b), 'b-', x, f1(x,a,b), 'b--', x, f2(x,a,b), 'g--', x, f3(x,a,b), 'r--', x, f4(x,a,b), 'k--', 'LineWidth', 3);
m = max(lik(x,a,b)) * 1.05;
line([a a], [0 m], 'Color', 'k', 'LineStyle', ':');
line([b b], [0 m], 'Color', 'k', 'LineStyle', ':');
legend('Overall likelihood', 'Z=(0,1)', 'Z=(1,0)', 'Z=(1,1)', 'Z=(0,0)', 'Location', 'NW')
axis tight
print -depsc lik.eps
pause

% assignment probabilities
p1 = phi(x-a) ./ (phi(x-a)+phi(0-a));
p2 = phi(x-b) ./ (phi(x-b)+phi(0-b));
plot(x, p1, x, p2, 'LineWidth', 3);
legend('P(Z_1=1)', 'P(Z_2=1)', 'Location', 'SE');
axis([min(x) max(x) -.05 1.05]);
axis tight
print -depsc probs.eps
pause

% EM iteration
mu = 0;
for iter = 1:10
    p1 = phi(mu-a) ./ (phi(mu-a)+phi(0-a));
    p2 = phi(mu-b) ./ (phi(mu-b)+phi(0-b));
    q = p1 .* (1-p2) .* log(f1(x,a,b)) + (1-p1) .* p2 .* log(f2(x,a,b));
    q = q + p1 .* p2 .* log(f3(x,a,b)) + (1-p1) .* (1-p2) .* log(f4(x,a,b));
    q = q - H(p1) - H(p2);
    [val, idx] = max(q);
    newmu = x(idx);
    % could also show following bound in log scale, demonstrating it is quadratic
    plot(x, lik(x,a,b), 'b-', x, exp(q), 'g-', 'LineWidth', 3);
    line([a a], [0 m], 'Color', 'k', 'LineStyle', ':');
    line([b b], [0 m], 'Color', 'k', 'LineStyle', ':');
    line([newmu newmu], [0 m], 'Color', 'g', 'LineStyle', '--');
    legend('Likelihood', sprintf('EM bound, mu=%.2f', mu), 'Location', 'NW');
    axis tight
    pause
    print('-depsc', sprintf('em%02d.eps', iter));
    mu = newmu;
end

end


% likelihood for mixture of 2 gaussians
function p = lik(x,a,b)
    p = f1(x,a,b)+f2(x,a,b)+f3(x,a,b)+f4(x,a,b);
end

function p = f1(x,a,b)
    p = phi(x-a) .* phi(0-b) .* .25;
end

function p = f2(x,a,b)
    p = phi(x-b) .* phi(0-a) .* .25;
end

function p = f3(x,a,b)
    p = phi(x-a) .* phi(x-b) .* .25;
end

function p = f4(x,a,b)
    z = zeros(size(x));
    p = phi(z-a) .* phi(z-b) .* .25;
end

% standard normal PDF
function p = phi(x) 
    p = (1/sqrt(2*3.141592654)) * exp(-.5*x.^2);
end

% entropy
function v = H(p)
    v = (p .* log(p) + (1-p) .* log(1-p));
end