/* external function declarations */

vector sort_rows( matrix _w, vector _s);
vector sort_cols( matrix _w, vector _s);

double get_cost( double _zz, int _cur, int _next, int _j, int _perm, int _el);	/* for K == 1 */
vector get_ex( vector _index);

vector get_perm( double _zz, int _vp, int _cur, int _next, int _j, int _el);	/* for K > 1 */
matrix get_paths( vector _index, matrix _vp, int _el);

void kalman_init(void);						/* for Kalman cost function */
void kalman_update( int _j, matrix _v, matrix _vl);
void kalman_perm_update( int _j, matrix _v, matrix _vl, matrix _vp);
vector get_kalman_perm( int _L, double _zz, int _vp, int _cur, int _next, int _j, int _el);

/* vit() */

void vit( int _t_size, int _tr_min)
{
/*
 * generic viterbi routine
 *
 * Finds L best of K completely unmerged paths.
 *
 * R. Perry, Oct. 1998
 *
 * References:
 *
 *   G. Forney, "The Viterbi Algorithm", Proceedings of the IEEE, Vol. 61,
 *   No. 3, March 1973, pp. 268-278.
 *
 *   J. Wolf, A. Viterbi, and G. Dixon, "Finding the Best Set of K Paths
 *   Through a Trellis with Application to Multitarget Tracking", IEEE
 *   Trans. on Aerospace and Electronic Systems, Vil. AES-25, No. 2, March
 *   1989, pp. 287-296.
 *
 *   N. Seshadri and C. Sundberg, "List Viterbi Decoding Algorithms with
 *   Applications", IEEE Trans. on Communications, Vol. 42, No. 2/3/4,
 *   Feb./March/April 1994, pp. 313-323.
 *
 * If K == 1, requires functions: get_cost(), get_ex()
 * If K >  1, requires functions: get_perm(), get_paths()
 */

/*
 * Skip declarations of most local variables here.
 * Let them remain global for help in debugging.
 *
 * matrix v;
 * vector good, pre, w, w_new, zz, tmp;
 * int NN, new_NN, t_size, tr, tr_min, i_min;
 * double w_min, p, perm;
 * 
 * vector path_index, ex;
 */

int i, j, el, j_start, k;

int t_default;	/* default trellis size */
string bell;	/* beep for warning messages */

VIT_NEW = 1;

t_default = 100;
bell = (string) 7;

/*** get data samples ***/

/*
 * get_x() must be called before calling vit()
 *
 * get_x() will create global variables:
 *
 *	N, M, K, states, zstate, x, y, noise, z
 *
 * states, x, y, and noise are not used in the algorithm itself, just
 * in some print() statements, and x is used at the end to check
 * the error between x and ex (estimated x).
 */

/* print( x, y, noise, z);	/* z is the input signal for the viterbi algorithm */

/*
 * Make temporary copy of z in zz,
 * since zz may be shifted if truncation of survivors occurs.
 */
zz = z;


/*** viturbi algorithm setup ***/

if( narg < 1)			/* set trellis size */
  t_size = t_default;
else if( _t_size > 1)
  t_size = _t_size;
else {
  print( bell, "\nvit: t_size argument is too small!  Resetting to t_default\n");
  t_size = t_default;
} 

NN = min( t_size, N);		/* set max trellis size */

if( narg < 2)
  tr_min = (int)(NN/2);		/* set minimum truncation length */
else
  tr_min = _tr_min;

if( tr_min > NN) {
  print( bell, "\nvit: tr_min is too big!  Resetting to NN/2\n");
  tr_min = (int)(NN/2);
} else if( tr_min < 1 ) {
  print( bell, "\nvit: tr_min is too small!  Resetting to NN/2\n");
  tr_min = (int)(NN/2);
}

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

L_prev = reshape( 1:L, 1, LM);	/* 1..L, M times */
I_prev = reshape( M1, L, M)';	/* 1..M, L times */
I_prev = I_prev(:);		/* e.g. for M=3, L=2 it is  1 1 2 2 3 3 */

vl = v = -ones( M*(NN+1), L);	/* viterbi path arrays */

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

if( zstate > 0) {		/* initial state is known ... */
  v(zstate,1) = zstate;
  vl(zstate,1) = 1;
  w(zstate,1) = 0;
} else {			/* or initial state is unknown */
  v(M1,1) = M1';
  vl(M1,1) = ones(M,1);
  w(M1,1) = 0;
}

if( K > 1) {			/* path permutation array */
  vp = v;
  vp(M1,1) = ones(M,1);
}

RP = 1;
if( K > 1 && L > 1)
  RP = rows(perms);

ex = zeros(LK,1);		/* we will append output values to ex later */


/* for kalman cost function */

    if( undef(KALMAN))
      KALMAN=1;

    if( KALMAN) {
      print( "Using Kalman cost function.");
      kalman_init();
    } else
      print( "Using non-Kalman cost function.");


/*** main viturbi loop ***/

j_start = 1;

while( 1) {	/* loop breaks at end when input data is exhausted */

  print( "-------\nj_start = ", j_start);

  for( j = j_start; j <= NN; ++j) {	/* for each column, i.e. time index */
  
    if( DEBUG > 1) print( "j = ", j);

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

    for( k = 1; k <= M; ++k) {		/* for each next state */
  
      if( DEBUG > 1) print( "k = ", k);
  
      w_ii = HUGE_VAL*ones(LM,RP);

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

      for( ii = 1;  ii <= LM;  ++ii) { 	/* for each previous state */
  
	i = I_prev(ii);		/* state index */
        el = L_prev(ii);	/* "l" value */

        if( w(i,el) == HUGE_VAL)  /* unreachable state */
          continue;
  
        if( DEBUG > 1) print( "i = ", i, "el = ", el);
  
        if( K == 1) {

	  p = get_cost( zz(j), i, k, j, 1, el);	/* j is relative after truncation ... */
				     /* 1 there is fake permutation number */

          if( DEBUG > 1) print( p);

	} else { /* K > 1 */
  
	  if( KALMAN)
	    tmp = get_kalman_perm( L, zz(j), vp(i+(j-1)*M,el), i, k, j, el);
	  else
	    tmp = get_perm( L, zz(j), vp(i+(j-1)*M,el), i, k, j, el);

	  p = tmp(1,:);
	  perm_ii(ii,:) = tmp(2,:);

          if( DEBUG > 1) print( tmp);
	}
  
	w_ii(ii,:) = p + w(i,el);
	w_ii(ii,find(p==HUGE_VAL)) = HUGE_VAL;

      } /* end loop on ii */

      /* keep L best paths */

      ii_sort = sort( w_ii);
      ii_sort = ii_sort(1:L);
      RR = sort_rows(w_ii,ii_sort);
      v(k+j*M,:) = I_prev(RR);
      vl(k+j*M,:) = L_prev(RR);
      w_ii = w_ii(:);
      w_new(k,:) = w_ii(ii_sort);
      v(k+j*M,find(w_new(k,:)==HUGE_VAL)) = -1;

      if( K > 1) {
	perm_ii = perm_ii(:);
	vp(k+j*M,:) = perm_ii(ii_sort);
      }

      if( DEBUG > 1) {
	print( w_new, v, vl);
	if( K > 1)
	  print( vp);
      }

    } /* end loop on k */
  
    w = w_new;

    if(KALMAN)
      if( K > 1)
        kalman_perm_update( j, v(M1+j*M,:), vl(M1+j*M,:), vp(M1+j*M,:));
      else
        kalman_update( j, v(M1+j*M,:), vl(M1+j*M,:));
  
    if( DEBUG == 1) print( j, v, vl, w);

  } /* end loop on j */
  
  
  /*** check optimal paths ***/
  
  ii_min = sort( w);

  if( DEBUG) print( ii_min);

  ii_min = ii_min( 1:L);

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

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

  print( w, i_min, l_min);
  
  /* check truncation */
  
  tr = 0;
  v_cur = v(1+NN*M:(NN+1)*M,:);
  vl_cur = vl(1+NN*M:(NN+1)*M,:);

  pre = v_cur(:);
  pre = pre(find(pre != -1));       /* indices of previous states */

  el_index = path_index = zeros(L,NN);

  if( K > 1) {
    vp_cur = vp(1+NN*M:(NN+1)*M,:);
    perm_index = zeros(L,NN);
  }

  for( j = NN; j >= 1; --j) {   /* create path_index and set truncation index */

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

    ii_min = ll_min = zeros(1,L);

    if( K > 1)
      pp_min = zeros(1,L);

    for( i = 1;  i <= L;  ++i) {
      ii_min(i) = v_cur( i_min(i), l_min(i));
      ll_min(i) = vl_cur( i_min(i), l_min(i));
      if( K > 1)
        pp_min(i) = vp_cur( i_min(i), l_min(i));
    }

    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,:);
    }

    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 != -1));
      }
    }

  }

  print( tr);

  if( DEBUG) print( path_index);

  if( cols(zz) == NN) {	/* used all of the input data */
    tr = NN;
    print( "input data finished");
  }
  
  if( tr < tr_min) {
    print( "\nvit: tr < tr_min, truncating at tr_min\n");
    tr = tr_min;
  }

  path_index = path_index(:,1:tr);

  if( K == 1)
    ex = [ex get_ex( path_index)];
  else {
    ex_tmp = zeros( LK, cols(path_index));
    for( el = 1; el <= L; ++el) {
      r_tmp = get_paths( path_index, perm_index, el);	/* intermediate r_tmp variable should not be needed */
      ex_tmp( 1+(el-1)*K:el*K, :) = r_tmp;		/* bug in m ??? */
    }
    ex = [ex ex_tmp];
  }
  
  if( cols(zz) == NN)		/* Done.  Input data is exhausted. */
    break;
  
  zz = zz(tr+1:cols(zz));	/* shift input data */

  if( DEBUG) print(zz);

  j_start = NN + 1 - tr;	/* next start column in the trellis */

  /* reset v and NN */

/*** fix the following section later... ***/

  new_NN = min( NN, cols(zz));
  v = [v(:,tr+1:NN+1), -ones(M,new_NN-NN+tr)];
  if( K > 1)
    vp = [vp(:,tr+1:NN+1), -ones(M,new_NN-NN+tr)];
  NN = new_NN;

  tr_min = (int)(NN/2);		/* reset minimum truncation length */

  good = find( w != HUGE_VAL);	/* renormalize w */
  w(good) = w(good) - w_min;

}
/*** end main viterbi loop ***/
  

print( "-------");

ex = ex(:,2:cols(ex));	/* remove extraneous initial column */

/* display results */
/*
print( ex, x);		

if( K == 1)
  print( "norm(ex-x) = ", norm(ex-x));
*/

}
