
function [P_y,P_xj_given_y,loglik] = discrete_em(niter,X,N)

% N is number of possible values of y = 1,...,N
X=X+1; % elements of X are now 1 or 2 (more convenient for indexing)
n = size(X,1); % number of samples
d = size(X,2); % number components in each sample

% initialize model P
P_y = ones(N,1)/N; % uniform pmf over y
P_xj_given_y = cell(d,1); % store one transition prob matrix P(xj|y) for each j=1,...,d
for j=1:d  % indexes components of x
    P = rand(2,N); % transition probabilities from y=1,...,N to xk=0,1
    P = P * diag(sum(P).^(-1)); % normalize columns to sum to one.
    P_xj_given_y{j} = P;  % conditional pmf for xj given y
end

% allocate storage for Q(y|xi), y=0,1, i=1,...,n.
Q = zeros(2,n);

% run EM algorithm...
loglik = zeros(niter,1);
for k = 1:niter

    % E-step: compute Q(y|x^i) for y=0,1 and i=1,...,n from P(y|x)
    for i=1:n % iterate over samples
        x = X(i,:);
        P_x = 0.0;
        for y=1:N  % calc prob of y given sample i from P
            Q(y,i)=P_y(y);
            for j=1:d
                Q(y,i)=Q(y,i)*P_xj_given_y{j}(x(j),y);
            end
            P_x = P_x + Q(y,i);   
        end         
        loglik(k) = loglik(k) + log(P_x);
        Q(:,i) = Q(:,i)/P_x; % normalize column to sum to one       
    end              
    loglik(k) = loglik(k)/n;  % avg log-lik per sample    
    fprintf('iter: %d, loglik: %f\n',k,loglik(k));

    % M-step: update P based on Q
    for y=1:N
        n_y = sum(Q(y,:)');
        P_y(y) = n_y/n;
        for j=1:d
            for xj=1:2
                n_y_given_xj = sum(Q(y,find(X(:,j)==xj)));
		P_xj_given_y{j}(xj,y) = n_y_given_xj / n_y;
            end
        end
    end

end
