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

/* ngram_decode.c --- for Julius: predict next word using N-gram probability
                      and word trellis index (for 2nd pass)*/

/* $Id: ngram_decode.c,v 1.6 2002/09/11 22:02:33 ri Exp $ */

/* N-gramを用いたスタックデコーディング(第2パス)における次単語集合の決定:
   与えられた展開元仮説の予測される始端フレーム周辺の単語トレリス上に
   残っている単語について，そのN-gram確率とともに返す */

/* Prediction of next words on N-gram based stack decoding (2nd pass):
   survived words in word trellis around the estimated end frame of source
   hypothesis will be returned with their N-gram probabilities. */

/* Julius モードでは ngram_firstwords(), ngram_nextwords(),
   ngram_acceptable(), ngram_eosscore() が第2パスのメイン関数 wchmm_fbs()
   に渡され, 使用される．(Julian モードでは "dfa_decode.c" 内の dfa_*() が
   使われる */

/* In julian mode, pointers to functions ngram_firstwords(), ngram_nextwords(),
   ngram_acceptable() and ngram_eosscore() are passed to main search function
   wchmm_fbs().  (for julian, corresponding functions dfa_*() in
   "dfa_decode.c" will be used instead) */
   
#include <julius.h>

#ifdef USE_NGRAM

/* for sorting next word candidates by their word ID */
static int
compare_nw(NEXTWORD **a, NEXTWORD **b)
{
  if ((*a)->id > (*b)->id) return 1;
  if ((*a)->id < (*b)->id) return -1;
  return 0;
}
/* find next word candiate whose id 'w' */
static NEXTWORD *
search_nw(NEXTWORD **nw, WORD_ID w, int num)
{
  int left,right,mid;
  NEXTWORD *tmp;

  if (num == 0) return NULL;
  left = 0;
  right = num - 1;
  while (left < right) {
    mid = (left + right) / 2;
    if ((nw[mid])->id < w) {
      left = mid + 1;
    } else {
      right = mid;
    }
  }
  tmp = nw[left];
  if (tmp->id == w) {
    return tmp;
  } else {
    return NULL;
  }
}

/*********** transparent words handling on 2nd pass **********/
static WORD_ID cnword[2];	/* current context words */
static int cnnum;		/* num of found context words */
static int last_trans;		/* num of skipped transparent words */
/* From cseq[0..n-1], set the last context words (= non-transparent words in
   hypothesis context) to cnword[0] and cnword[1].
   Also, the num of skipped transparent words are
   counted in last_trans.
 */
static void
set_word_context(WORD_ID *cseq, int n, WORD_INFO *winfo)
{
  WORD_ID cw[2];
  int i;

  cnnum = 0;
  last_trans = 0;
  for(i=n-1;i>=0;i--) {
    if (! winfo->is_transparent[cseq[i]]) {
      cnword[cnnum++] = winfo->wton[cseq[i]];
      if (cnnum >= 2) break;
    } else {
      last_trans++;
    }
  }
}

/*************************************************************/

/* lookup survived words on specified frame in backtrellis, 
   compute their N-gram probabilities, and add them to NEXTWORD data */
static int
pick_backtrellis_words(
		       BACKTRELLIS *bt,	/* backtrellis info */
		       WORD_INFO *winfo, /* word dictionary */
		       NGRAM_INFO *ngram, /* N-gram info */
		       NEXTWORD **nw, /* next words info (return value) */
		       int oldnum, /* initial offset of nw */
		       NODE *hypo, /* source hypothesis */
		       short t	/* time frame */
		       )
{
  int i;
  WORD_ID w;
  int nword1=0, nword2=0;
  LOGPROB rawscore;
#ifdef WPAIR
  int w_old = WORD_INVALID;
#endif
  int num;

  num = oldnum;
  /* set word context to cnword[0], cnword[1] */
  set_word_context(hypo->seq, hypo->seqnum, winfo);
  /* lookup survived words in backtrellis on time frame 't' */
  for (i=0;i<bt->num[t];i++) {
    w = (bt->rw[t][i])->wid;
#ifdef WORD_GRAPH
    /* only words on the word graphs are expanded */
    if (!(bt->rw[t][i])->within_wordgraph) continue;
    /* word connection is restricted by the graph structure */
    if (w != ((hypo->tre)->last_tre)->wid) continue;
#else /* not WORD_GRAPH */
#ifdef WPAIR
    /* some word have same word ID with different previous word, so
       only one will be opened (best word will be selected later
       by next_word() */
    if (w == w_old) continue;	/* backtrellis is sorted by word ID */
    else w_old = w;
#endif /* WPAIR */
#endif /* not WORD_GRAPH */
    /* skip if already exist */
    if (search_nw(nw, w, oldnum) != NULL) continue;
    switch(cnnum) {		/* length of valid context */
    case 0:			/* unigram */
#ifdef CLASS_NGRAM
      rawscore = class_uni_prob(ngram, winfo->wton[w]);
#else
      rawscore = uni_prob(ngram, winfo->wton[w]);
#endif
      break;
    case 1:			/* bigram */
#ifdef CLASS_NGRAM
      rawscore = class_bi_prob_rl(ngram, winfo->wton[w], cnword[0]);
#else
      rawscore = bi_prob_rl(ngram, winfo->wton[w], cnword[0]);
#endif
      break;
    default:			/* trigram */
#ifdef CLASS_NGRAM
      rawscore = class_tri_prob_rl(ngram, winfo->wton[w], cnword[0], cnword[1]);
#else
      rawscore = tri_prob_rl(ngram, winfo->wton[w], cnword[0], cnword[1]);
#endif
      break;
    }
    nw[num]->tre   = bt->rw[t][i];
    nw[num]->id    = w;
    nw[num]->lscore = lm_weight2 * rawscore + lm_penalty2;
    if (winfo->is_transparent[w]) {
      /*nw[num]->lscore -= (LOGPROB)last_trans * TRANS_RENZOKU_PENALTY;*/
      if (winfo->is_transparent[hypo->seq[hypo->seqnum-1]]) {
	nw[num]->lscore += lm_penalty_trans;
      }
    }
    
    /* j_printf("%d: %s added\n", num, winfo->wname[nw[num]->id]); */
    num++;
  }
  return num;
}

/* Look for survived backtrellis words near the specified frame, and
   make NEXTWORD list.
   Words in frame [tm-lookup_range..tm+lookup_range-1] will be picked up.
   If a word exist in several frames, only one near the center frame
   will be taken: the true connection point will be determined later at
   next_word() */
int
get_backtrellis_words(
		      BACKTRELLIS *bt,	/* backtrellis info */
		      WORD_INFO *winfo, /* word dictionary */
		      NGRAM_INFO *ngram, /* N-gram info */
		      NEXTWORD **nw, /* next words info (return value) */
		      NODE *hypo, /* source hypothesis */
		      short tm,	/* center frame to lookup */
		      short t_end /* end frame of input */
		      )
{
  int num = 0;
  int t, t_step;
  int oldnum=0;

  if (tm < 0) return(0);

#ifdef PREFER_CENTER_ON_TRELLIS_LOOKUP
  /* fix for 3.2 (01/10/18 by ri) */
  /* before and after (one near center frame has high priority) */
  for (t_step = 0; t_step < lookup_range; t_step++) {
    /* before or center */
    t = tm - t_step;
    if (t < 0 || t > bt->framelen - 1 || t >= t_end) continue;
    num = pick_backtrellis_words(bt, winfo, ngram, nw, oldnum, hypo, t);
    if (num > oldnum) {
      qsort(nw, num, sizeof(NEXTWORD *),
	    (int (*)(const void *,const void *))compare_nw);
      oldnum = num;
    }
    if (t_step == 0) continue;	/* center */
    /* after */
    t = tm + t_step;
    if (t < 0 || t > bt->framelen - 1 || t >= t_end) continue;
    num = pick_backtrellis_words(bt, winfo, ngram, nw, oldnum, hypo, t);
    if (num > oldnum) {
      qsort(nw, num, sizeof(NEXTWORD *),
	    (int (*)(const void *,const void *))compare_nw);
      oldnum = num;
    }
  }

#else

  /* before the center frame */
  for(t = tm; t >= tm - lookup_range; t--) {
    if (t < 0) break;
    num = pick_backtrellis_words(bt, winfo, ngram, nw, oldnum, hypo, t);
    if (num > oldnum) {
      qsort(nw, num, sizeof(NEXTWORD *),
	    (int (*)(const void *,const void *))compare_nw);
      oldnum = num;
    }
  }
  /* after the center frame */
  for(t = tm + 1; t < tm + lookup_range; t++) {
    if (t > bt->framelen - 1) break;
    if (t >= t_end) break;
    num = pick_backtrellis_words(bt, winfo, ngram, nw, oldnum, hypo, t);
    if (num > oldnum) {
      qsort(nw, num, sizeof(NEXTWORD *),
	    (int (*)(const void *,const void *))compare_nw);
      oldnum = num;
    }
  }
#endif

  return num;
}

/* check limitation of next word */
/* return value: nextword num */
int
limit_nw(NEXTWORD **nw, NODE *hypo, int num)
{
  int src,dst;
  int newnum;

  /* <s>からは何も展開しない */
  /* no hypothesis will be generated after "<s>" */
  if (hypo->seq[hypo->seqnum-1] == winfo->head_silwid) {
    return(0);
  }

  dst = 0;
  for (src=0; src<num; src++) {
    if (nw[src]->id == winfo->tail_silwid) {
      /* </s> は展開しない */
      /* do not expand </s> (it only appears at start) */
      continue;
    }
    if (src != dst) memcpy(nw[dst], nw[src], sizeof(NEXTWORD));
    dst++;
  }
  newnum = dst;

  return newnum;
}
	


/* 最初の単語群を返す．返り値: 単語数 (-1 on error) */
/* return initial word set.  return value: num of words (-1 on error) */
int
ngram_firstwords(
		 NEXTWORD **nw,	/* next word set (return value) */
				/* (assume already malloced) */
		 int peseqlen,	/* end time of input */
		 int maxnw,	/* maximum length of nw[] */
		 WORD_INFO *winfo, /* word dictionary */
		 BACKTRELLIS *bt)	/* backtrellis info */
{
#ifdef WORD_GRAPH
  int last_time;
#endif

#ifdef SP_BREAK_CURRENT_FRAME
  if (rest_param != NULL) {
    /* 初期仮説は 最終フレームに残った単語トレリス上の最尤単語 */
    /* the initial hypothesis is the best word survived on the last frame of
       the segment */
    nw[0]->id = sp_break_2_begin_word;
  } else {
    /* 最終セグメント: 初期仮説は 単語の末尾の無音単語(=winfo->tail_silwid) */
    /* we are in the last of sentence: initial hypothesis is word-end silence word */
    nw[0]->id = winfo->tail_silwid;
  }
#else
  /* initial hypothesis is word-end silence word */
  nw[0]->id = winfo->tail_silwid;
#endif
#ifdef FIX_PENALTY
  nw[0]->lscore = 0.0;
#else
  nw[0]->lscore = lm_penalty2;
#endif
#ifdef WORD_GRAPH
  /* 初期仮説のTRELLIS_ATOMをここで決定する */
  /* in word-graph mode, the first trellis atom should be defined here */
  for (last_time = peseqlen - 1; last_time >= 0; last_time--) {
    nw[0]->tre = bt_binsearch_atom(&backtrellis, last_time, winfo->tail_silwid);
    if (nw[0]->tre != NULL) break;
  }
  if (nw[0]->tre == NULL) {		/* </s> not found */
    j_printerr("no wordend left in beam!\n");
    j_printerr("beam width too small or input speech too short\n");
    return 0;
  }
#endif

  return 1;			/* number of words = 1 */
}

/* ある仮説の次の接続単語群を返す．帰り値: 単語数 (-1 on error) */
/* return next word set from the hypothesis.  return value:
   num of words (-1 on error) */
int
ngram_nextwords(
		NODE *hypo,
		NEXTWORD **nw,
		int maxnw,	/* hypo: source */
		NGRAM_INFO *ngram, /* N-gram info */
		WORD_INFO *winfo, /* word dictionary */
		BACKTRELLIS *bt) /* backtrellis info */
{
  int num, num2;

  if (hypo->seqnum == 0) {
    j_error("gs_get_next_words: hypo contains no word\n");
  }

  /* 仮説の推定終端時刻において backtrellis内に残っている単語を得る */
  /* get survived words on backtrellis at the estimated end frame */
  num = get_backtrellis_words(bt, winfo, ngram, nw, hypo, hypo->estimated_next_t, hypo->bestt);

  if (debug2_flag) j_printf("%d",num);

  /* 展開できない単語をチェックして外す */
  /* exclude unallowed words */
  num2 = limit_nw(nw, hypo, num);

  if (debug2_flag) j_printf("-%d=%d unfolded\n",num-num2,num2);

  return(num2);
}

/* 仮説が文として受理可能であるかどうかを返す */
/* return if the hypothesis is "acceptable" */
boolean
ngram_acceptable(NODE *hypo, WORD_INFO *winfo)
{
  if (
#ifdef SP_BREAK_CURRENT_FRAME
      /* 最後の仮説が第１パス最尤仮説の最初の単語と一致しなければならない */
      /* the last word should be equal to the first word on the best hypothesis on 1st pass */
      hypo->seq[hypo->seqnum-1] == sp_break_2_end_word
#else
      /* 最後の仮説が文頭無音単語でなければならない */
      /* the last word should be head silence word */
      hypo->seq[hypo->seqnum-1] == winfo->head_silwid
#endif
      ) {
    return TRUE;
  } else {
    return FALSE;
  }
}

#endif /* USE_NGRAM */ /* ------------------------------------------ */
