/* data function declarations */

double get_yval( int _cur, int _next);		    /* for K == 1 */
vector get_ex( vector _index);

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

/* vit() */

void vit( int _t_size, int _tr_min)
{
/*
 * generic viterbi routine
 *
 * Finds best K completely unmerged paths.
 *
 * References:
 *
 *   G. D. Forney Jr., The Viterbi Algorithm,
 *   Proceedings of the IEEE, Vol. 61, No. 3,
 *   March 1973, pp. 268-278
 *
 *   J. K. Wolf, A. M. Viterbi, and G. S. Dixon,
 *   Finding the Best Set of K Paths Through a Trellis
 *   with Application to Multitarget Tracking,
 *   IEEE AES-25, No. 2, March 1989, pp. 287-296
 *
 * If K == 1, requires functions: get_yval(), get_ex(), and pn()
 * 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, yval, perm;
 * 
 * vector path_index, ex;
 */

int i, j, j_start, k;
int t_default;	/* default trellis size */
string bell;	/* beep for warning messages */

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

v = -ones( M, NN+1);		/* viterbi path array */
if( zstate > 0)
  v( zstate, 1) = zstate;
else
  v(:,1) = [1:M]';

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

if( zstate > 0) {
  w = HUGE_VAL*ones( 1, M);	/* path length measures */
  w( zstate) = 0;		/* initial state */
} else
  w = zeros( 1, M);

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


/*** 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 */
  
    w_new = HUGE_VAL*ones( 1, M);	/* set big values to start with */
  
    for( i = 1; i <= M; ++i) {		/* for each row, i.e. state index */
  
      if( w(i) == HUGE_VAL)		/* HUGE_VAL marks unreachable states */
        continue;
  
      if( DEBUG > 1) print( "i = ", i);
  
      for( k = 1; k <= M; ++k) {	/* for each possible next state */
  
        if( K == 1) {

          yval = get_yval( i, k, j);	/* corresponding y value */
					/* some examples only use i and k args */

          if( yval == HUGE_VAL)		/* unreachable state */
	    continue;
  
          p = pn( zz(j) - yval);	/* =~ -ln Pn(z-y) */

          if( DEBUG > 1) print( yval, p);

	} else {			/* K > 1 */
  
	  tmp = get_perm( zz(j), vp, i, k, j);
	  p = tmp(1);
	  perm = tmp(2);

          if( DEBUG > 1) print( perm, p);

	  if( p == HUGE_VAL)		/* unreachable state */
	    continue;
	}
  
        if( p + w(i) < w_new(k)) {	/* found better path */
          w_new(k) = p + w(i);
          v( k, j+1) = i;
	  if( K > 1)
	    vp( k, j+1) = perm;
        }
      }
      if( DEBUG > 1) {
	print( w_new, v);
	if( K > 1)
	  print( vp);
      }
    }
  
    w = w_new;
  
    if( DEBUG > 1) print( j, v, w);
  }
  
  
  /*** check optimal path ***/
  
  i_min = find( w == min(w));
  i_min = i_min(1);
  w_min = w(i_min);
  
  print( i_min, w);
  
  /* check truncation */
  
  tr = 0;
  pre = find(v(:,NN+1) != -1);	/* indices of previous states */
  path_index = zeros(1,NN);
  
  for( j = NN; j >= 1; --j) {	/* create path_index and set truncation index */
  
    path_index(j) = i_min;
    i_min = v( i_min, j+1);
  
    if( tr == 0) {
      if( all( pre == pre(1))) { /* found node which all survivors go through */
        tr = j;
      } else {
        pre = find(v( pre, j) != -1);
      }
    }
  }
      
  print( tr);

  if( DEBUG) print( "[0 path_index]", [0 path_index], v, "col #", [1:NN+1]);

  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 = [ex get_paths( path_index, vp)];
  
  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 */

  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));
*/

}
