% 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 w0;	% initial costs
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 MP;	% number of permutations

global merge_period;	% how often to merge similar paths
global tr_period;	% how often to truncate the trellis

global NBEST;		% for alternative N-best pruning

% 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_save;		% saved costs for post-processing
global w_min;		% final minimum costs
global path_index;	% final state paths
global el_index;	%  and el indices
global perm_index;	%  and permutation indices

% for window K estimation
%
global Wtot;
global WINDOW;		% flag for windowing
global Winpath_index;	% Holds final state paths for the estimating window
global Winel_index;	% Same for el_index
global Winperm_index;	%  and perm_index
global windowNum;	% Counts what window its calculating
global estMar;		% #tracks marginalization estimate
global WinMM;		% subset of MM for the current window
global MM;

global MM_max;
global Winc;
global XY;
global kxv;
global combos;
global missed_cost;

windowNum = 1;

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

w_save = zeros( L, N*M);	% saved costs for post-processing

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);
  if( zstate < 0)		% initial state is known statistically
    w(1,M1) = w0;  		% with initial cost vector w0
  else
    w(1,M1) = 0;
  end
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 = MP;			% size(perms,2);
end

perm_ii = 0;

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

path_index = zeros(L,N);	% initialize path_index, el_index, and perm_index
el_index = path_index;
if( K > 1)
  perm_index = zeros(L,N);
end

merge_count = 0;		% initialize merge and tr counters
tr_count = 0;

tr = 0;				% tr is time index of previous trellis truncation

%%% main viterbi loop %%%

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

  % j

  merge_count = merge_count + 1;

  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

    % path merging
    %
    % if( (merge_count == merge_period) | ((j == N) & (merge_period ~= HUGE_VAL)) )
    if( merge_count == merge_period)

      if( DEBUG) disp('path merging...'), k, end;

      w_ii = vit_merge( tr, j, k, w_ii, I_prev, L_prev, perm_ii);

    end % path merging

    % 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;

%%%
%%
% for N-best comparison, Radar2000

  if( NBEST)

    w = w(:);

    [tmp ii_min] = sort(w);

    ii_min = ii_min(1:L);

    w = HUGE_VAL*ones(LM,1);
    
    w(ii_min) = tmp(1:L);

    w = reshape(w,L,M);

  end
%
%%
%%%

  w_save(:,(j-1)*M+M1) = w;	% saved costs for post-processing

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

  % some cost functions (kalman) need to be updated

  jM = M1 + j*M;

  if( K > 1)
    perm_update( j, v(:,jM), vl(:,jM), vp(:,jM));
  else
    cost_update( j, v(:,jM), vl(:,jM));
  end

  if( DEBUG == 1) j, end;

  if( DEBUG > 1) v, vl, w, end;

  % check/reset merge_count
  %
  if( merge_count == merge_period)
    merge_count = 0;
  end


  % trellis truncation
  %
  tr_count = tr_count + 1;

  if( (tr_count == tr_period) | (j == N))

    tr_count = 0;

    if( DEBUG) disp('trellis truncation...'), end;

    [i_min l_min] = vit_tr( tr, j);

    tr = j;

% ------------------------------- %
%
% Window K estimation, from Jess
%
    if( WINDOW & tr_period ~= HUGE_VAL)

      disp ('Estimating Tracks...')

      % Save the window to Winpath, Winel, and Winperm indices
      %
      Winpath_index = path_index(:,j-tr_period+1:j);
      Winel_index   = el_index(:,j-tr_period+1:j);
      if( K > 1)
        Winperm_index   = perm_index(:,j-tr_period+1:j);
      end

      WinMM = MM(j-tr_period+1:j);

      % Perform window estimation
      %
      PKKwindow;

      % Record answer
      %
      recordEst;
      windowNum = windowNum + 1;

    %%%
    %%
    % trying this part, 22 March 2000, RP
    % to handle truncation in the middle of the trellis
    %
    % update w and Wtot: keep the L-best and set the rest to HUGE_VAL
    %
    % Note: this will interfere with PKK!
    %
%      w_tmp = HUGE_VAL*ones( L, M);
%      w_tmp(l_min,i_min) = w(l_min,i_min);
%      w = w_tmp;
%      Wtot = zeros( K*M, L);
%
      if( j < N)

        w = HUGE_VAL*ones( L, M);
        w(1,M1) = 0;

        % similar to cost_init():

        i1 = j*M*K + [1:M*K];               % index for next kxv values
	kxv(:,i1) = 0;

        txy = zeros( 2, max(M,MM(j)));        % max(M,MM(j)) handles K==MM(j) case
        txy(:,1:MM(j)) = XY(:,1:MM(j));
        for i = 1:M
          i1 = 1 + j*M*K + (i-1)*K;
          i1 = [i1:i1+K-1];
          kxv(1,i1) = txy(1,combos(:,i));
          kxv(3,i1) = txy(2,combos(:,i));
        end
  
        Winc = zeros( L*K*M, MM_max+1);
        Winc(:, MM_max+1) = missed_cost;
        Wtot = zeros( K*M, L);
  
        kalman_predict(j+1);

      end

    %
    %%
    %%%

    end % if WINDOW

% ------------------------------- %

    % path_index, el_index

  end % trellis truncation

end % j

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

if( DEBUG > 1) path_index, el_index, end;
