% -----------------------------------
% EM estimation of the hmm parameters

function [hmm,lik] = esthmm(X,hmm);

lik = []; cont = 1; it=0;
while (cont),
  [hmm,lx] = EMstep(X,hmm);
  lik = [lik;lx];
  if (it>5), cont = (lx-lik(it)>0.0001); end;
  it = it + 1;

  if (mod(it,10)==0), disp(lx); end;
end;


% ------------------------------
% one EM iteration

function [hmm,lik] = EMstep(X,hmm)

[L,n] = size(X); 

[ns,no] = size(hmm.Po);

lik = 0;

Ns0 = zeros(1,ns);  % soft counts for the starting state 
Nss = zeros(ns,ns); % soft counts for state transitions
Nsx = zeros(ns,no); % soft counts for state/observation pairs

% the E-step

for l=1:L,
  x = X(l,:);

  a = alpha(x,hmm); % forward propagation
  b = beta(x,hmm);  % backward propagation

  Px = sum(a(:,n)); % prob of observed sequence

  g = (a.*b)/Px;       % posterior state probs
  Ns0 = Ns0 + g(:,1)'; 

  xt = zeros(no,n);     % indicator function representation of 
  xt(x+no*(0:n-1)) = 1; % observations (e.g., [0;1])

  Nsx = Nsx + g*xt';

  bx  = hmm.Po(:,x).*b;
  Nss = Nss + (a(:,1:n-1)*bx(:,2:n)').*hmm.Pt/Px; % transition counts per seq
                                                  
  lik = lik + log(Px); % log probability of data 
end;

% the M-step

hmm.Pt = normrows(Nss);
hmm.Po = normrows(Nsx);
hmm.p0 = normrows(Ns0);

