/* Copyright (c) 1991-2002 Doshita Lab. Speech Group, Kyoto University */
/* Copyright (c) 2000-2002 Speech and Acoustics Processing Lab., NAIST */
/*   All rights reserved   */

/* vsegment.c --- generic routine to do viterbi segmentation */

/* $Id: vsegment.c,v 1.5 2002/11/08 13:44:51 ri Exp $ */

/* any segmentatino unit is allowed: segmentation point should be specified by sequence of HMM number to pickup */

#include <sent/stddefs.h>
#include <sent/htk_param.h>
#include <sent/hmm.h>

typedef struct __seg_token__ {
  int last_id;
  int last_end_frame;
  LOGPROB last_end_score;
  struct __seg_token__ *next;
  struct __seg_token__ *chain;
} SEGTOKEN;

LOGPROB				/* returns the acoustic score */
viterbi_segment(HMM *hmm,	/* concatinated sentence HMMs */
	     HTK_Param *param,	/* input parameter vectors */
	     int *endstates,	/* where unit ends in the hmm */
	     int ulen,		/* total unit numin the hmm */
	     int **id_ret,       /* return sequence of result */
	     int **seg_ret,	/* return segmented frame number at each endstates */
	     LOGPROB **uscore_ret,/* return normalized scores for each unit */
	     int *slen_ret)
{
  /* for viterbi */
  LOGPROB *nodescore[2];	/* node buffer */
  SEGTOKEN **tokenp[2];		/* propagating token which holds segment info */
  int *from_node;
  int i, n, t;
  int tl,tn;
  LOGPROB tmpsum;
  A_CELL *ac;
  SEGTOKEN *newtoken, *token, *tmptoken;
  LOGPROB result_score;
  LOGPROB maxscore, minscore;	/* for debug */
  int maxnode;			/* for debug */
  SEGTOKEN *chain_root;
  int *id, *seg, slen;
  LOGPROB *uscore;

#if 0
  for(i=0;i<hmm->len;i++) {
    printf("%d ->", i);
    for (ac = hmm->state[i].ac; ac; ac = ac->next) {
      printf(" %d (%f)", ac->arc, ac->a);
    }
    printf("\n");
  }
#endif

  /* assume more than 1 units */
  if (ulen < 1) {
    j_printerr("Error: viterbi_segment: no unit?\n");
    return LOG_ZERO;
  }

  /* initialize node buffers */
  tn = 0;
  tl = 1;
  for (i=0;i<2;i++){
    nodescore[i] = (LOGPROB *)mymalloc(hmm->len * sizeof(LOGPROB));
    tokenp[i] = (SEGTOKEN **)mymalloc(hmm->len * sizeof(SEGTOKEN *));
  }
  from_node = (int *)mymalloc(sizeof(int) * hmm->len);

  chain_root = NULL;
  for (n = 0; n < hmm->len; n++) {
    nodescore[tn][n] = LOG_ZERO;
    newtoken = (SEGTOKEN *)mymalloc(sizeof(SEGTOKEN));
    newtoken->last_id = -1;
    newtoken->last_end_frame = -1;
    newtoken->last_end_score = 0.0;
    newtoken->next = NULL;
    tokenp[tn][n] = newtoken;
    newtoken->chain = chain_root;
    chain_root = newtoken;
  }
  /* first frame: only set initial score */
  /*if (hmm->state[0].is_pseudo_state) {
    j_printerr("Warning: state %d: pseudo state?\n", 0);
    }*/
  nodescore[tn][0] = 0.0;

  /* do viterbi for frames */
  for (t = 0; t <= param->samplenum; t++) {
    i = tl;
    tl = tn;
    tn = i;
    maxscore = LOG_ZERO;
    minscore = 0.0;

    /* clear next scores */
    for (i=0;i<hmm->len;i++) {
      nodescore[tn][i] = LOG_ZERO;
      from_node[i] = -1;
    }

    /* select viterbi path for each node */
    for (n = 0; n < hmm->len; n++) {
      if (nodescore[tl][n] <= LOG_ZERO) continue;
      for (ac = hmm->state[n].ac; ac; ac = ac->next) {
        tmpsum = nodescore[tl][n] + ac->a;
        if (nodescore[tn][ac->arc] < tmpsum) {
          nodescore[tn][ac->arc] = tmpsum;
	  from_node[ac->arc] = n;
	}
      }
    }
    /* propagate token, appending new if path was selected between units */
    for (n = 0; n < hmm->len; n++) {
      if (from_node[n] == -1 || nodescore[tn][n] <= LOG_ZERO) {
	/*tokenp[tn][n] = NULL;*/
      } else {
	i=0;
	while (from_node[n] > endstates[i]) i++;
	if (n > endstates[i]) {
	  newtoken = (SEGTOKEN *)mymalloc(sizeof(SEGTOKEN));
	  newtoken->last_id = i;
	  newtoken->last_end_frame = t-1;
	  newtoken->last_end_score = nodescore[tl][from_node[n]];
	  newtoken->next = tokenp[tl][from_node[n]];
	  tokenp[tn][n] = newtoken;
	  newtoken->chain = chain_root;
	  chain_root = newtoken;
	} else {
	  tokenp[tn][n] = tokenp[tl][from_node[n]];
	}
      }
    }

    /* if this is next of last frame, loop ends here */
    if (t == param->samplenum) break;
    
    /* calc outprob to new nodes */
    for (n = 0; n < hmm->len; n++) {
      if (hmm->state[n].out.state == NULL) continue;
      if (nodescore[tn][n] > LOG_ZERO) {
	if (hmm->state[n].is_pseudo_state) {
	  j_printerr("Warning: state %d: pseudo state?\n", n);
	}
	nodescore[tn][n] += outprob(t, &(hmm->state[n]), param);
      }
      if (nodescore[tn][n] > maxscore) { /* for debug */
	maxscore = nodescore[tn][n];
	maxnode = n;
      }
    }
    
#if 0
    for (i=0;i<ulen;i++) {
      printf("%d: unit %d(%d-%d): begin_frame = %d\n", t - 1, i,
	     (i > 0) ? endstates[i-1]+1 : 0, endstates[i],
	     (tokenp[tl][endstates[i]] == NULL) ? -1 :
	     tokenp[tl][endstates[i]]->last_end_frame + 1);
    }
#endif

    /* printf("t=%3d max=%f n=%d\n",t,maxscore, maxnode); */
    
  }

  result_score = nodescore[tn][hmm->len-1];

  /* parse back the last token to see the trail of best viterbi path */
  /* and store the informations to returning buffer */
  slen = 0;
  for(token = tokenp[tn][hmm->len-1]; token; token = token->next) {
    if (token->last_end_frame == -1) break;
    slen++;
  }
  id = (int *)mymalloc(sizeof(int)*slen);
  seg = (int *)mymalloc(sizeof(int)*slen);
  uscore = (LOGPROB *)mymalloc(sizeof(LOGPROB)*slen);

  i = slen - 1;
  for(token = tokenp[tn][hmm->len-1]; token; token = token->next) {
    if (i < 0 || token->last_end_frame == -1) break;
    id[i] = token->last_id;
    seg[i] = token->last_end_frame;
    uscore[i] = token->last_end_score;
    i--;
  }

  /* normalize scores by frame */
  for (i=slen-1;i>0;i--) {
    uscore[i] = (uscore[i] - uscore[i-1]) / (seg[i] - seg[i-1]);
  }
  uscore[0] = uscore[0] / (seg[0] + 1);

  /* set return value */
  *id_ret = id;
  *seg_ret = seg;
  *uscore_ret = uscore;
  *slen_ret = slen;

  /* free memory */
  free(from_node);
  token = chain_root;
  while (token) {
    tmptoken = token->chain;
    free(token);
    token = tmptoken;
  }
  for (i=0;i<2;i++) {
    free(nodescore[i]);
    free(tokenp[i]);
  }

  return(result_score);

}
