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

/* gprune_heu.c --- calculate probability of Gaussian densities */
/*                   with Gaussian pruning (heuristic) */

/* $Id: gprune_heu.c,v 1.2 2002/09/11 22:01:50 ri Exp $ */

#include <sent/stddefs.h>
#include <sent/htk_hmm.h>
#include <sent/htk_param.h>
#include <sent/hmm.h>
#include <sent/gprune.h>
#include "globalvars.h"

/* mail routines to calculate mixtures for given frame, and store to cache */
/*

  best_mixtures_on_last_frame[]
  
  dim:  0 1 2 3 4 .... veclen-1    -> sum up
 ================================
  backsum <-<-<-<-<-<-<-<-<-
 --------------------------------
  mix1  ->->
  mix2  ->->
  mix3  ->->
  ...
  mixN  ->->
 ================================
  score = forward sum + backmax

  algorithm 1:

     // store sub-optimal max prob for each dimension
     init backsum[]
     foreach all_mixtures in best_mixtures_on_last_frame {
        foreach dim {
           p = prob of this dim
	   if (backmax[dim] < p) backmax[dim] = p
	   sumup score
        }
	store to cache
     }
     // sum-up backward to make backmax
     foreach dim-1 backward {
        backmax[dim] += backmax[dim+1]
     }
     // compute rest with pruning with heuristics
     threshold = the current lowest score
     foreach rest_mixtures {
        foreach dim {
	   sumup score
	   if (score + backmax[dim+1] < threshold) {
	      skip this mixture
	   }
	}
	update threshold
     }
  }
     
*/

/* pruning threshold */
static LOGPROB *backmax;	/* backward sum of max for each dimension (inversed) */
static int backmax_num;		/* veclen */

static boolean *mixcalced;	/* mark which Gaussian has been computed */

/* initialize backmax */
static void
init_backmax()
{
  int i;
  for(i=0;i<backmax_num;i++) backmax[i] = 0;
}

/* sum-up backward to make backmax */
/*                        |
 *  0 1 2 ... max-3 max-2 | max-1
 *
 *  a b c      x      y   | ????
 *             |
 *             v
 *  .....     x+y     y     0.0
 */
static void
make_backmax()
{
  int i;
  backmax[backmax_num-1] = 0.0;
  /* backmax[len-1] = backmax[len-1];*/
  for(i=backmax_num-2;i>=0;i--) {
    backmax[i] += backmax[i+1];
  }
  /*  for(i=0;i<=len;i++) {
    printf("backmax[%d]=%f\n",i,backmax[i]);
    }*/
}

/* calculate probability while setting max values to backmax */
static LOGPROB
compute_g_heu_updating(HTK_HMM_Dens *binfo)
{
  VECT tmp, x, sum = 0.0;
  VECT *mean;
  VECT *var;
  VECT *bm = backmax;
  VECT *vec = OP_vec;
  short veclen = OP_veclen;

  if (binfo == NULL) return(LOG_ZERO);
  mean = binfo->mean;
  var = binfo->var->vec;

  tmp = 0.0;
  for (; veclen > 0; veclen--) {
    x = *(vec++) - *(mean++);
    tmp = x * x / *(var++);
    sum += tmp;
    if ( *bm < tmp) *bm = tmp;
    bm++;
  }
  return((sum + binfo->gconst) / -2.0);
}
/* calculate probability with pruning with backmax heuristics */
static LOGPROB
compute_g_heu_pruning(HTK_HMM_Dens *binfo, LOGPROB thres)
{
  VECT tmp, x;
  VECT *mean;
  VECT *var;
  VECT *bm = backmax;
  VECT *vec = OP_vec;
  short veclen = OP_veclen;
  LOGPROB fthres;

  if (binfo == NULL) return(LOG_ZERO);
  mean = binfo->mean;
  var = binfo->var->vec;
  fthres = thres * (-2.0);

  tmp = 0.0;
  bm++;
  for (; veclen > 0; veclen--) {
    x = *(vec++) - *(mean++);
    tmp += x * x / *(var++);
    if ( tmp + *bm > fthres) {
      return LOG_ZERO;
    }
    bm++;
  }
  return((tmp + binfo->gconst) / -2.0);
}


/* init */
boolean
gprune_heu_init()
{
  int i;
  /* maximum Gaussian set size = maximum mixture size */
  OP_calced_maxnum = OP_hmminfo->maxmixturenum;
  OP_calced_score = (LOGPROB *)mymalloc(sizeof(LOGPROB) * OP_gprune_num);
  OP_calced_id = (int *)mymalloc(sizeof(int) * OP_gprune_num);
  mixcalced = (boolean *)mymalloc(sizeof(int) * OP_calced_maxnum);
  for(i=0;i<OP_calced_maxnum;i++) mixcalced[i] = FALSE;
  backmax_num = OP_hmminfo->opt.vec_size + 1;
  backmax = (LOGPROB *)mymalloc(sizeof(LOGPROB) * backmax_num);

  return TRUE;
}

/* compute a set of Gaussians with safe pruning */
void
gprune_heu(HTK_HMM_Dens **g, int gnum, int *last_id)
{
  int i, j, num = 0;
  LOGPROB score, thres;

  if (last_id != NULL) {	/* compute them first to form thresholds */
    /* 1. clear backmax */
    init_backmax();
    /* 2. calculate first $OP_gprune_num with setting max for each dimension */
    for (j=0; j<OP_gprune_num; j++) {
      i = last_id[j];
      score = compute_g_heu_updating(g[i]);
      num = cache_push(i, score, num);
      mixcalced[i] = TRUE;      /* mark them as calculated */
    }
    /* 3. set backmax for each dimension */
    make_backmax();
    /* 4. calculate the rest with pruning*/
    thres = OP_calced_score[num-1];
    for (i = 0; i < gnum; i++) {
      /* skip calced ones in 1. */
      if (mixcalced[i]) {
        mixcalced[i] = FALSE;
        continue;
      }
      /* compute with safe pruning */
      score = compute_g_heu_pruning(g[i], thres);
      if (score > LOG_ZERO) {
	num = cache_push(i, score, num);
	thres = OP_calced_score[num-1];
      }
    }
  } else {			/* in case the last_id not available */
    /* at the first 0 frame */
    /* calculate with safe pruning */
    thres = LOG_ZERO;
    for (i = 0; i < gnum; i++) {
      if (num < OP_gprune_num) {
	score = compute_g_base(g[i]);
      } else {
	score = compute_g_safe(g[i], thres);
	if (score <= thres) continue;
      }
      num = cache_push(i, score, num);
      thres = OP_calced_score[num-1];
    }
  }
  OP_calced_num = num;
}
