%rng(123123123);

% Simulate data
theta_true = 4;
tau_true = 0.3;
n_samples = 10000;
z = (rand(n_samples,1) < tau_true) + 1; % 2 with probability tau_true
x = zeros(n_samples,1);
for i = 1:n_samples
    if z(i)==1
        x(i) = randn; % N(0,1)
    elseif z(i)==2
        x(i) = randn + theta_true;  % N(theta_true,1)
    end
end


% Parameters of the prior distributions.
alpha0 = 0.5;
beta0 = 0.2;

n_iter = 20;
% To keep track of the estimates of tau and theta in different iterations:
tau_est = zeros(n_iter,1);
th_est = zeros(n_iter,1);

% Some initial value for the things that will be updated
E_log_tau = -0.7;   % E(log(tau))
E_log_tau_c = -0.7;  % E(log(1-tau))
E_log_var = 4 .* ones(n_samples,1); % E((x_n-theta)^2)
r2 = 0.5 .* ones(n_samples,1); % Responsibilities of the second cluster.


for iter = 1:n_iter
    
    % Updates of responsibilites, factor q(z)
    log_rho1 = E_log_tau_c - 0.5 .* log(2*pi) - 0.5 .* (x.^2);
    log_rho2 = E_log_tau - 0.5 .* log(2*pi) - 0.5 .* E_log_var;
    max_log_rho = max(log_rho1, log_rho2); % Normalize to avoid numerical problems when exponentiating.
    rho1 = exp(log_rho1 - max_log_rho);
    rho2 = exp(log_rho2 - max_log_rho);
    r2 = rho2 ./ (rho1 + rho2);
    r1 = 1 - rho2;
    
    N1 = sum(r1);
    N2 = sum(r2);
    
    % Update of factor q(tau)
    E_log_tau = % EXERCISE
    E_log_tau_c = % EXERCISE
    
    % Update of factor q(theta)
    E_log_var = % EXERCISE
    
    % Keep track of the current estimates
    tau_est(iter) = (N2 + alpha0) / (N1 + N2 + 2*alpha0);
    th_est(iter) = m2;    
end

disp(num2str([tau_est th_est]));
% With large n_samples, this should converge to the (tau_true, theta_true).