% generic viterbi routine
%
% Finds L best of K completely unmerged paths.
%
% Functions called in here, to be supplied by the user for
% a particular problem type:
%
%   cost_init()
%
%	To initialize the cost function parameters.
%
%	This is a noop for ISI problems.  For Kalman tracking
%	problems cost_init() should initialize the Kalman filters
%	and gain parameters.
%
%   For K=1:
%
%	get_cost()
%
%		To get the cost value for a state transition.
%
%	cost_update()
%
%		To allow the cost function to update itself
%		based on the Viterbi decisions.  This is a
%		noop for ISI, but necessary for Kalman.
%
%   For K>1:
%
%	get_perm()
%
%		To get the cost value and permutation index for a state transition.
%
%	perm_update()
%
%		To allow the cost function to update itself
%		based on the Viterbi decisions.  This is a
%		noop for ISI, but necessary for Kalman.
%
% R. Perry, Jan. 1999

% comment out the following function line to make this a script,
% then local variables will be accessible for debugging.
%
% function vit()

% general globals
%
global HUGE_VAL;
global DEBUG;

% global inputs
%
global z;	% input data row vector
global zstate;	% initial state value
global K;	% number of tracks
global L;	% number of best to keep
global M;	% number of states
global N;	% number of input data points
global perms;	% permutation array

% global outputs
%
global v;		% main viterbi state array
global vl;		% array for L-best indices
global vp;		% array for permutation indices
global w;		% cost array
global w_min;		% final minimum costs
global i_min;		%  and state indices
global l_min;		%  and el indices
global path_index;	% final state paths
global el_index;	%  and el indices
global perm_index;	%  and permutation indices

M1=[1:M];
M2 = M*M;
LM = L*M;
LK = L*K;

L_prev = [1:L]' * ones(1,M);
L_prev = L_prev(:);		% 1..L, M times

I_prev = ones(L,1) * M1;	% 1..M, L times
I_prev = I_prev(:);             % e.g. for M=3, L=2 it is  1 1 2 2 3 3

v = zeros( L, M*(N+1));		% viterbi path arrays
vl = v;

w = HUGE_VAL*ones( L, M);	% cost array


if( zstate > 0) 		% initial state is known
  v(1,zstate) = zstate;
  vl(1,zstate) = 1;
  w(1,zstate) = 0;
else 				% initial state is unknown
  v(1,M1) = M1;
  vl(1,M1) = ones(1,M);
  w(1,M1) = 0;
end

if( K > 1) 			% path permutation array
  vp = v;
  vp(1,M1) = ones(1,M);
end


RP = 1;				% number of rows for w_ii and perm_ii
if( (K > 1) & (L > 1))
  RP = size(perms,2);
end

cost_init;			% some cost functions (kalman) need to be initialized

%%% main viterbi loop %%%

for j = 1:N			% for each column, i.e. time index

  if( DEBUG > 1) j, end;

  w_new = HUGE_VAL*ones(L,M);	% set big values to start with 

  for k = M1			% for each next state 1..M

    if( DEBUG > 1) k, end;

    w_ii = HUGE_VAL*ones(RP,LM);

    if( K > 1)
	perm_ii = ones(RP,LM);
    end

    for ii = 1:LM		% for each previous state
  
      i = I_prev(ii);		% state index
      el = L_prev(ii);		% "l" value

      if( w(el,i) ~= HUGE_VAL)	% reachable state

        if( DEBUG > 1) i, el, end;

        if( K == 1)
	  p = get_cost( z(j), i, k, j, el);
          if( DEBUG > 1) p, end;
	else
	  [p perm_ii(:,ii)] = get_perm( L, z(j), vp(el,i+(j-1)*M), i, k, j, el);
          if( DEBUG > 1) p, perm_ii(:,ii), end;
	end

	w_ii(:,ii) = p + w(el,i);
	if( any(p==HUGE_VAL))
	  w_ii(find(p==HUGE_VAL),ii) = HUGE_VAL;
	end

      end % reachable state

    end % ii

    % keep L best paths

    [tmp ii_sort] = sort( w_ii(:));
    ii_sort = reshape( ii_sort(1:L), L, 1);
    RR = sort_cols(w_ii,ii_sort);

    v(:,k+j*M) = I_prev(RR);
    vl(:,k+j*M) = L_prev(RR);

    w_ii = w_ii(:);
    w_ii = reshape( w_ii(ii_sort), L, 1);
    w_new(:,k) = w_ii;
    if( any(w_ii==HUGE_VAL))
      v(find(w_ii==HUGE_VAL),k+j*M) = 0;
    end

    if( K > 1)
      perm_ii = perm_ii(:);
      perm_ii = reshape( perm_ii(ii_sort), L, 1);
      vp(:,k+j*M) = perm_ii;
    end

    if( DEBUG > 1)
      w_new, v, vl
      if( K > 1) vp, end;
    end

  end % k

  w = w_new;

  if( all( w(:) == HUGE_VAL))
    error( 'Viterbi is stuck.  All states are unreachable.')
  end

  % some cost functions (kalman) need to be updated

  if( K > 1)
    perm_update( j, v(:,M1+j*M), vl(:,M1+j*M), vp(:,M1+j*M));
  else
    cost_update( j, v(:,M1+j*M), vl(:,M1+j*M));
  end

  if( DEBUG == 1) j, v, vl, w, end;

end % j


%%% check optimal paths %%%
  
[tmp ii_min] = sort(w(:));

if( DEBUG) ii_min, end;

ii_min = ii_min(1:L);

i_min = sort_cols(w,ii_min);
l_min = sort_rows(w,ii_min);	% should eliminate HUGE_VAL paths here ???

w_min = w(:);
w_min = w_min(ii_min);

% w, i_min, l_min

% check truncation 

tr = 0;
v_cur = v(:,1+N*M:(N+1)*M);
vl_cur = vl(:,1+N*M:(N+1)*M);

pre = v_cur(:);
pre = pre(find(pre ~= 0));       % indices of previous states

path_index = zeros(L,N);
el_index = path_index;

if( K > 1)
  vp_cur = vp(:,1+N*M:(N+1)*M);
  perm_index = zeros(L,N);
end

for j = N:-1:1	% create path_index and set truncation index

  path_index(:,j) = i_min;
  el_index(:,j) = l_min;

  ll_min = zeros(L,1);
  ii_min = ll_min;

  if( K > 1)
    pp_min = ll_min;
  end

  for i = 1:L
    ii_min(i) = v_cur( l_min(i), i_min(i));
    ll_min(i) = vl_cur( l_min(i), i_min(i));
    if( K > 1)
      pp_min(i) = vp_cur( l_min(i), i_min(i));
    end
  end

  i_min = ii_min;
  l_min = ll_min;

  v_cur = v(:,1+(j-1)*M:j*M);
  vl_cur = vl(:,1+(j-1)*M:j*M);

  if( K > 1)
    perm_index(:,j) = pp_min;
    vp_cur = vp(:,1+(j-1)*M:j*M);
  end

  % following might not be right for L>1; also need to look at vl ???

  if( tr == 0)
    if( all( pre == pre(1))) % found node which all survivors go through
      tr = j;
    else
      pre = v_cur(:);
      pre = pre(find(pre ~= 0));
    end
  end

end % j

path_index = path_index';
el_index = el_index';
if( K > 1)
  perm_index = perm_index';
end


% tr

if( DEBUG) path_index, el_index, end;
