%Author: Ashish Kapoor MIT Media Lab, 10/20/02
%Please send corrections to: ash@media.mit.edu
function [pi_out, a_out, b_out, llk_array]=learn_hmm_param(pi_in, a_in, b_in, num_states, obs_data, EPSILON);
%Learn parameters of HMM until abs(newloglikelihood-oldloglikelihood) <= EPSILON%

pi_out = pi_in;
a_out = a_in;
b_out = b_in;

%Memory Allocation

K = size(obs_data,2);
alpha_hat = cell(K, 1);
alpha_hat_hat = cell(K, 1);
c=cell(K, 1);
beta_hat = cell(K, 1);
beta_hat_hat = cell(K, 1);

llk=0; old_llk=-999;
llk_array=[];

while abs(llk-old_llk) > EPSILON
    pi=pi_out;
    a=a_out;
    b=b_out;
    
    %**********************************************************************************
    %Compute alphas and betas corresponding to each observation

    for k=1:K
        %Initialization
        obs=obs_data{k};
        T = size(obs,1);
    
        alpha_hat{k} = zeros(T, num_states);
        alpha_hat_hat{k} = zeros(T, num_states);
        c{k} = zeros(T,1);
        beta_hat{k} = zeros(T, num_states);
        beta_hat_hat{k} = zeros(T, num_states);
    
        %Computation of alphas
        %Init alphas and c
        alpha_hat_hat{k}(1,:)=pi'.*b(:,obs(1)+1)';
        c{k}(1)=1/sum(alpha_hat_hat{k}(1,:));
        alpha_hat{k}(1,:)=c{k}(1)*alpha_hat_hat{k}(1,:);
        %Recursion for alphas
        for t=2:T
            alpha_hat_hat{k}(t,:)=(alpha_hat{k}(t-1,:)*a(:,:)).*b(:,obs(t)+1)';
            c{k}(t)=1/sum(alpha_hat_hat{k}(t,:));        
            alpha_hat{k}(t,:)=c{k}(t)*alpha_hat_hat{k}(t,:);    
        end
    
        %Computation of betas
        %Init betas
        beta_hat_hat{k}(T,:)=1;
        beta_hat{k}(T,:)=c{k}(T)*beta_hat_hat{k}(T,:);
        %Recursion for betas
        for t=1:T-1
            beta_hat_hat{k}(T-t,:)=a(:,:) * (b(:,obs(T-t+1)+1) .* beta_hat{k}(T-t+1,:)') ;
            beta_hat{k}(T-t,:)=c{k}(T-t)*beta_hat_hat{k}(T-t,:);      
        end
    end

    %**********************************************************************************
    %Update pi, a and b
   
    %pi stays the same
    pi_out=pi;

    %Update a and b
    for i=1:num_states
        num_a=zeros(1,size(a,2));
        num_b=zeros(1,size(b,2)); 
        denom_a=0;
        denom_b=0;
            
        for k=1:K
            obs=obs_data{k};
            T = size(obs,1);

            for t=1:T-1
                num_a = num_a + alpha_hat{k}(t,i)*a(i,:).*b(:,obs(t+1)+1)'.*beta_hat{k}(t+1,:);
                num_b(obs(t)+1) = num_b(obs(t)+1) + (alpha_hat{k}(t,i)*beta_hat{k}(t,i)/c{k}(t));
            end
            
            denom_a=denom_a + sum((alpha_hat{k}(1:T-1,i).*beta_hat{k}(1:T-1,i)./c{k}(1:T-1)));
            denom_b=denom_b + sum((alpha_hat{k}(1:T,i).*beta_hat{k}(1:T,i)./c{k}(1:T)));
            num_b(obs(T)+1) = num_b(obs(T)+1) + (alpha_hat{k}(T,i)*beta_hat{k}(T,i)/c{k}(T));
        end
    
        a_out(i,:)=num_a/denom_a;
        b_out(i,:)=num_b/denom_b;
    end
    
    %***********************************************************************************
    %Compute Convergence Criterea
    llk = 0;
    old_llk = 0;

    for k=1:K
        llk = llk + log_prob_obs_hmm(pi_out, a_out, b_out, num_states, obs_data{k});
        old_llk = old_llk + (-sum(log(c{k})));
    end
    
    %Plot loglikelihoods
    if size(llk_array)==0
        fprintf(1, 'Initial Loglikelihood = %f\n', old_llk);
        llk_array=[old_llk llk];
    else
        llk_array=[llk_array llk];
    end

    fprintf(1, 'Current Loglikelihood = %f\n', llk);
end