function [] = HMM_EM_demo(N_state,tau,scale_mu,scale_cov,N_count),

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%                       HMM_EM	         	      %      
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% HMM_EM_demo(N_state,tau,scale_mu,scale_cov,N_count)%
%%					              %
%% N_state = number of state [3]	              %
%% tau = lentgh of the sequence	[200]	              %
%% scale_mu = scale of the mean [4-5]	              %
%% scale_cov = scale of the cov [1]	              %
%% N_count = number of iteration [100]	              %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%
%% Please refer to  Max Welling's class notes 
%% for the symbol notation.
%% In figure (1) there is the estimated state evolution
%% using the Viterbi algorithm given A, mu and B  and the 
%% sequence of observations (with A, mu and B selected 
%% randomly). 
%% In figure(2) there is the sequence of observations (x's) 
%% The color is the label of the state attached to the observation
%% In figure(3) there is the estimated state evolution computed 
%% using the Viterbi algorithm. In this case A, mu and B 
%% are estimated using  the EM algorithm over the observations.
%% The red  plots in fig.2 are the real ni and lambda 
%% The white plots in fig.2 are the estimated ni and lambda 
%% The squares in fig.2 are the incorrectly estimated states. 
%% If there are too many mistakes the relabeling procedure fails 
%% and the squares are no meaningful any longer.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%    S. savarese   M.Welling -   May 2000
%%%%    Vision Lab  - Caltech 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%% X = observation vector
%% Y = state vector

if (nargin==0)
   N_state = 3; %% number of states
   N_count = 100;
   tau = 100;    %% length of the state vector
   scale_mu = 5;
   scale_cov = 1;
end;


colordef black;
%clear;
%close all;

figure(1), clf;
figure(2), clf;

figure(2),
title ('green=state 1; red=state 2; magenta=state 3; yell=state 4; green=state 5; ...')

%%%%%%%% initialization

color = ['g' 'r' 'm' 'y' 'c' 'b' 'w' ];
%N_state = 3; %% number of states
N_obs   = 1;
%N_count = 100;
%tau = 200;    %% length of the state vector
dim_Y = 2;   %% dimension of the state  
dim_X = 2;   %% dimension of observation

%scale_mu = 4;
%scale_cov = 1;
X = zeros(dim_X,tau); %%observation
Y =  zeros(1,tau);   %% state
Y_est =  zeros(1,tau);   %% state estimate
Y_mu  = zeros(dim_Y,N_state);
Y_cov = zeros(dim_Y,dim_Y,N_state);
Y_mu_new  = zeros(dim_Y,N_state);
Y_cov_new = zeros(dim_Y,dim_Y,N_state);
Y_axis1 = 1:1:tau;

for i=1:N_state,
    Y_mu(:,i)=randMean(dim_Y,scale_mu);
    Y_cov(:,:,i)=randCovariance(dim_Y,scale_cov);
end;

gamma_i    = zeros(N_state,tau);
omega_ij   = zeros(N_state,N_state,tau);
A_ij	   = zeros(N_state,N_state);
B_i	   = zeros(N_state,tau);
B_i_new	   = zeros(N_state,tau);
A_ij_new   = zeros(N_state,N_state);
B_i_new	   = zeros(N_state,tau);
alpha_i  = zeros(N_state,tau);
kt       = zeros(1,tau);
beta_i   = zeros(N_state,tau);
mu_i     = zeros(N_state,tau);
mu_i_new = zeros(N_state,1);
delta_i	 = zeros(N_state,tau);
ni_i_new	= zeros(dim_X,N_state);
lambda_i_new    = zeros(dim_X,dim_X,N_state);

%%%%%% generate transition probability matrix

A1_ij	 = zeros(N_state,N_state);
A1_ij(:,:) = rand(N_state,N_state);
A2=sum(A1_ij,1);
A3=ones(N_state,1)*A2(1,:);
A_ij=A1_ij./A3;
    
%%%%%%%% generate P(Yo=i) vector

P2_Yo	    = rand(N_state,1);
P3_Yo =sum(P2_Yo,1)*ones(N_state,1);
P_Yo = P2_Yo./P3_Yo;


%%%%%%%%% compute the initial state

Yo = new_state(P_Yo,1); %% label of the state
Y(1)=Yo;
    

%%%%%% generate the sequence of states

for i=2:tau,
    Y(i) = new_state(A_ij,Y(i-1));
end;


%%%%%%% plot the sequence of state

figure(1);
axis([0 tau+1 0 N_state+1]);
grid on;
title('square = actual state');
drawnow;


%%%%% generate the sequence of observations

for i=1:tau,    
 
   X(:,i)= randvec(1,Y_mu(:,Y(i)),Y_cov(:,:,Y(i)));
    figure(2);
    hold on;
    plot(X(1,i),X(2,i),[color(Y(i)) 'x'],'linewidth',2);
    figure(1);
    hold on
    plot(Y_axis1(i),Y(i),[color(Y(i)) 's'],'linewidth',2);
    
    %if (i<10) pause(1); end;
    drawnow;
end;



%%%%%%%%% init parameters  for the Viterbi algorithm
%%%%%%%%% i.e. : mu_i(:,1)  B_i  delta_i(:,1)

for i=1:N_state,
    for t=1:tau,
	B_i(i,t)  = gauss2(Y_mu(:,i),Y_cov(:,:,i),X(:,t));
    end;
end;

mu_i(:,1) = P_Yo;

delta_i(:,1) = B_i(:,1).*mu_i(:,1);
delta_i(:,1)=delta_i(:,1)/norm(delta_i(:,1));



%%%%%%%%%%%% store values

Y_mu_o  = Y_mu;
Y_cov_o = Y_cov;
A_ij_o = A_ij;
mu_i_o = mu_i(:,1);
B_i_o = B_i;


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%% intial condition for the EM algorithm %%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%% Y_mu, Y_cov
for i=1:N_state,
    Y_mu_new(:,i)=randMean(dim_Y,scale_mu);
    Y_cov_new(:,:,i)=eye(dim_X);
end;

%%%%%%%%%%%%%%%% B_i
for i=1:N_state,
    for t=1:tau,
	B_i_new(i,t)  = gauss2(Y_mu_new(:,i),Y_cov_new(:,:,i),X(:,t));
    end;
end;


%%%%%%%%%%%%%%% mu_i
P2_Yo	    = rand(N_state,1);
P3_Yo =sum(P2_Yo,1)*ones(N_state,1);
P_Yo = P2_Yo./P3_Yo;
mu_i_new(:,1) = P_Yo;


%%%%%%%%%%%%%%%%% A_ij
%A1_ij	 = zeros(N_state,N_state);
A1_ij(:,:) = rand(N_state,N_state);
A2=sum(A1_ij,1);
A3=ones(N_state,1)*A2(1,:);
A_ij_new=A1_ij./A3;

A_err=A_ij-A_ij_o;
A_err=sqrt(A_err.^2);
err=sum(sum(A_err,1));

%%%%%

A_ij = A_ij_new;
mu_i = mu_i_new(:,1);
B_i = B_i_new;


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%         EM algorithm                  %%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
count = 1;

while (count<N_count),

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%  E-step   %%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%  Viterbi Algorithm

t1=zeros(N_state,1);
psi=zeros(N_state,tau);

for t=1:tau-1,
    for i=1:N_state,
	t1=A_ij(i,:).*delta_i(:,t)';
        [max_t1,index] = max(t1);
	delta_i(i,t+1) = B_i(i,t+1)*max_t1;
	psi(i,t+1)=index;
    end;

    delta_i(:,t+1)=delta_i(:,t+1)/norm(delta_i(:,t+1));
end;


%%%%%%%% compute alpha and kt

a1 = sum( B_i(:,1).*mu_i(:,1));
alpha_i(:,1) = B_i(:,1).*mu_i(:,1)/a1;

kt(1) = sum(B_i(:,1).*mu_i(:,1));
a2 = zeros(1,N_state);

for t=2:tau,
    for i=1:N_state,
	  a2(i) = (A_ij(i,:)*alpha_i(:,t-1))*B_i(i,t);
    end;
    kt(t)= sum(a2);
    alpha_i(:,t) = a2(:)/kt(t);
end; 
    


%%%%%%%%%%  back tracking

list=((-tau+1):-1);
list=list*(-1);

[max_t1 index] = max(delta_i(:,tau));
Y_est(tau) = index;

for t=list,
    delta_i(:,t);
    Y_est(t) = psi(Y_est(t+1),t+1);
end;



%%%%%%% compute beta

beta_i(:,tau) = ones(N_state,1);

for t=tau:-1:2,
 for j=1:N_state,
    beta_i(j,t-1)= sum(B_i(:,t).*beta_i(:,t).*A_ij(:,j))/kt(t);
 end;
end;


%%%%%%%%  compute gamma

for t=1:tau,
    gamma_i(:,t) = alpha_i(:,t).*beta_i(:,t);
end;


%%%%%%%% compute omega

for t=2:tau,
    beta_m = beta_i(:,t)*ones(1,N_state);
    B_m = B_i(:,t)*ones(1,N_state);
    alpha_m = ones(N_state,1)*alpha_i(:,t-1)';
    omega_ij(:,:,t) = (B_m.*beta_m.*A_ij.*alpha_m)/kt(t);
end;

%for t=2:tau,
%    for i=1:N_state,
%	for j=1:N_state, 
%	    omega_ij(i,j,t)=B_i(i,t)*beta_i(i,t)...
%		*A_ij(i,j)*alpha_i(j,t-1)/kt(t);
%
%	end;
%    end;
%end;



%%%%%%%%% compute error

%error = not((Y_est==Y));
%index2 = find(error);

%figure(2);
%hold on;
%plot(X(1,index2),X(2,index2),'wo');
%drawnow

%figure(1), hold off; ;
%figure(2), hold off;


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%  M-step   %%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%



mu_i_new = (1/N_obs)*gamma_i(:,1);

aij = ones(N_state,1)*sum(gamma_i(:,1:tau-1),2)';
A_ij_new = sum(omega_ij,3)./aij;

a=sum(gamma_i,2)';

%for i=1:N_state,
%   ni_i_new(:,i) = (sum((ones(2,1)*gamma_i(i,:)).*X,2)/a(i));
%end;

ni_t = zeros(dim_X,tau);

for i=1:N_state,
   for t=1:tau,
      ni_t(:,t) = gamma_i(i,t)*X(:,t);
   end;
   ni_i_new(:,i)=sum(ni_t,2)/a(i);
end;

lambda_t = zeros(dim_X,dim_X,tau);
a=sum(gamma_i,2)';

for i=1:N_state,
   for t=1:tau,
       lambda_t(:,:,t)=gamma_i(i,t)...
         *(X(:,t) - ni_i_new(:,i))*(X(:,t) - ni_i_new(:,i))';
   end; 
   lambda_i_new(:,:,i)=sum(lambda_t,3)/a(i)+ eye(2)*1E-5;
end;

%+ eye(2)*1E-5


for i=1:N_state,
    for t=1:tau,
	B_i_new(i,t)  = gauss2(ni_i_new(:,i),lambda_i_new(:,:,i),X(:,t));
    end;
end;

%%%%%%%%%%%%%% debug service  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
for i=1:N_state,
    for t=1:tau,
%%%%%%%%%%%%%%%  try to avoid huge peak
if (B_i_new(i,t)>100),
   fprintf('B_i = %f\n',B_i_new(i,t));
   ni_i_new(:,i)=randMean(dim_Y,scale_mu);
   lambda_i_new(:,:,i)=eye(2);
   plotGauss_color(Y_mu_new(1,i),Y_mu_new(2,i),Y_cov_new(1,1,i),Y_cov_new(2,2,i),Y_cov_new(2,1,i),'b');

   for i=1:N_state,
       for t=1:tau,
	B_i_new(i,t)  = gauss2(ni_i_new(:,i),lambda_i_new(:,:,i),X(:,t));
       end;
   end;

   fprintf('alterating loop...\n');
end;
%%%%%%%%%%%%%%  end try to avoid huge peak
end;
end;
%%%%%%%%%%%%%%  end debug service %%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%  updating
mu_i(:,1) = mu_i_new;
A_ij = A_ij_new;
B_i = B_i_new;

A_err=A_ij-A_ij_o;
A_err=sqrt(A_err.^2);
err=sum(sum(A_err,1));


%%%%% plot the sequence of observation
figure(2);
clf;

hold on;
plot(X(1,:),X(2,:),'yx','linewidth',2);
%if (i<10) pause(1); end;
drawnow;



if (count==1),   %% plot initial mu and cov
   %%%%%%%% plot Y_mu and Y_cov
   figure(2);
   hold on;
   for i=1:N_state
        plotGauss_color(Y_mu_new(1,i),Y_mu_new(2,i),Y_cov_new(1,1,i),Y_cov_new(2,2,i),Y_cov_new(2,1,i),'b');
   end;
   %input('');
end;


%%%%%%%% plot ni and sigma
figure(2);
hold on;
for i=1:N_state
    plotGauss_color(ni_i_new(1,i),ni_i_new(2,i),lambda_i_new(1,1,i),lambda_i_new(2,2,i),lambda_i_new(2,1,i),'w');
end;

%%%%%%%% plot Y_mu and Y_cov
figure(2);
hold on;
for i=1:N_state,
    plotGauss_color(Y_mu(1,i),Y_mu(2,i),Y_cov(1,1,i),Y_cov(2,2,i),Y_cov(2,1,i),'r');
end;




%%%%% loglikelihood

L(count)=log(prod(kt));
fprintf('L=%f\n',L(count));


if (count>1)&(sqrt((L(count)-L(count-1))^2)<0.00001),
  break;
end;

if (count>1)&(L(count)<L(count-1)),
   %fprintf('ERROR!!!\n');
   %return;
end;

count = count +1;
%input('press return');

end; %%%%%(while loop)


%%%%%%%%%%% relabel new  estimated state 
Rit = gamma_i;
[max_R Y_resp] = max(Rit,[],1);

y_i3=zeros(1,N_state);
Y_resp2=ones(1,tau)*(N_state+1);
A_ij2 = zeros(N_state,N_state);
ii = ones(1,N_state);

for i=1:N_state,
    for j=1:N_state,

	[index1] = find(Y==i);
	[y_i1] = (Y==i);
        [index2] = find(Y_resp==j);
	[y_i2] = (Y_resp==j);
	y_i3(j) = sum(y_i1.* y_i2);
   end;
   [maxi index3] = max(y_i3);
   j=index3;
   [index2] = find(Y_resp==j);
   Y_resp2(index2)=i;
   ii(i)=j;
end;

for i=1:N_state,
    for j=1:N_state,
	A_ij2(i,j)=A_ij(ii(i),ii(j));
    end;
end;

%%%%%%%%% compute error

error = not((Y_resp2==Y));
index2 = find(error);

%figure(2);
%hold on;
%plot(X(1,index2),X(2,index2),'wd');

%figure(1), hold off;
%figure(2), hold off;





%%%%% plot the sequence of observation
figure(2);
clf;
for i=1:tau,    
    hold on;
    plot(X(1,i),X(2,i),[color(Y(i)) 'x'],'linewidth',2);
    %if (i<10) pause(1); end;
    drawnow;
end;


%%%%%%%% plot ni and sigma
figure(2);
hold on;
for i=1:N_state
    plotGauss_color(ni_i_new(1,i),ni_i_new(2,i),lambda_i_new(1,1,i),lambda_i_new(2,2,i),lambda_i_new(2,1,i),'w');
end;

%%%%%%%% plot Y_mu and Y_cov
figure(2);
hold on;
for i=1:N_state,
    plotGauss_color(Y_mu(1,i),Y_mu(2,i),Y_cov(1,1,i),Y_cov(2,2,i),Y_cov(2,1,i),'r');
end;





%%%%%%% plot the sequence of state

figure(3);
clf;
axis([0 tau+1 0 N_state+1]);
grid on;
title('Square = actual state; yellow dot = estimated state');
drawnow;


%%%%% plot  the sequence of estimated states
for i=1:tau,
    figure(3);
    hold on
    plot(Y_axis1(i),Y_resp2(i),[color(Y_resp2(i)) 's'],'linewidth',2);
    %if (i<10) pause(1); end;
    grid on
    drawnow;
end;

%%%%%%% plot the real sequence of state 
figure(3);
hold on;
plot(Y_axis1,Y,'y.','linewidth',4);
axis([0 tau+1 0 N_state+1]);
grid on;




%%%%%% plot error on 2D graph
figure(2);
hold on;
plot(X(1,index2),X(2,index2),'ws');
drawnow;


A_ij_o = A_ij_o
A_ij2 = A_ij2

save obs2 X
save solution2 A_ij_o A_ij2 B_i_new mu_i_new