% kaleh_iter - one EM iteration for the Kaleh/Vallet algorithm
%
% Assumes (-1,1) symbol values.
%
% \bibitem{kaleh}
% G. K. Kaleh and R. Vallet,
% ``Joint Parameter Estimation and Symbol Detection for Linear or Nonlinear Unknown Channels,''
% {\em IEEE Trans. on Communications,} Vol. 42, No. 7, pp. 2406-2413, July 1994.
%
function kaleh_iter

global states;	% standard Viterbi globals
global N;	% number of samples
global M;	% number of states
global z;	% received data
global var_noise;
global FIR;
global FIR_length;

global PS;	% state transition probability matrix
global Xt;	% X(t) vector values
global Nij;	% used to get i and j given n
global IJn;	% used to get n given i and j
global gamma;   % array for forward recursion
global beta;    % array for backward recursion
global LY;      % array for likelihood values

global gamma_scale;	% normalization for columns of gamma and beta

M2 = 2*M;

%  Precompute py
%
py=zeros(M2,N);

for t=1:N
   for n=1:M2
      py(n,t) = exp(-(z(t)-FIR'*Xt(:,n))^2/(2*var_noise));
   end
end

% py = py-min(py(:));
% py = exp(-py);


% Compute gamma and beta
%
for t=2:N+1

   for j=1:M
      r = 0;
      for i=1:M
         n = IJn(i,j);
         if n ~= 0
	    % py = exp(-(z(t-1)-FIR'*Xt(:,n))^2/(2*var_noise));
            r = r + gamma(i,t-1)*PS(i,j)*py(n,t-1);
         end
      end
      gamma(j,t) = r;
   end

   % gamma(:,t) = gamma(:,t)/sum(gamma(:,t));

   gamma_scale(t) = sum(gamma(:,t));

   gamma(:,t) = gamma(:,t)/gamma_scale(t);

end 

for t=N:-1:1

    for i=1:M
      r = 0;
      for j=1:M
         n = IJn(i,j);
         if n ~= 0
            % py = exp(-(z(t)-FIR'*Xt(:,n))^2/(2*var_noise));
            r = r + beta(j,t+1)*PS(i,j)*py(n,t);
         end
      end
      beta(i,t) = r;
   end

   % beta(:,t) = beta(:,t)/sum(beta(:,t));

   beta(:,t) = beta(:,t)/gamma_scale(t);

end

% Compute LY
%
for t=1:N

   for n=1:M2
      ij = Nij(:,n);
      i = ij(1);
      j = ij(2);
      LY(n,t) = py(n,t)*PS(i,j)*gamma(i,t)*beta(j,t+1);
   end
   
   % LY(:,t) = LY(:,t)/sum(LY(:,t));

end

% Reestimate FIR
%
b = zeros(FIR_length,1);
for n=1:M2
   b = b + (LY(n,:)*z')*Xt(:,n);
end

A = zeros(FIR_length,FIR_length);
for n=1:M2
   A = A + (sum(LY(n,:))*Xt(:,n))*Xt(:,n)';
end

% if A is singular to working precision, A\b can result
% in FIR -> Inf, then the Viterbi algorithm will fail.
%
%  FIR = A\b;
%
% safer solution:
%
  FIR = pinv(A)*b;
