#!/usr/bin/env bash

# Copyright 2012  Johns Hopkins University (Author: Guoguo Chen)
# Apache 2.0

# Begin configuration section.
model= # You can specify the model to use
cmd=run.pl
acwt=0.083333
lmwt=1.0
max_silence_frames=50
max_states=1000000
max_states_scale=4
max_expand=180 # limit memory blowup in lattice-align-words
strict=true
word_ins_penalty=0
silence_word=  # Specify this only if you did so in kws_setup
skip_optimization=false     # If you only search for few thousands of keywords, you probablly
                            # can skip the optimization; but if you're going to search for
                            # millions of keywords, you'd better do set this optimization to
                            # false and do the optimization on the final index.
frame_subsampling_factor=   # We will try to autodetect this. You should specify
                            # the right value if your directory structure is
                            # non-standard
# End configuration section.

echo "$0 $@"  # Print the command line for logging

[ -f ./path.sh ] && . ./path.sh; # source the path.
. parse_options.sh || exit 1;

if [ $# != 4 ]; then
   echo "Usage: steps/make_index.sh [options] <kws-data-dir> <lang-dir> <decode-dir> <kws-dir>"
   echo "... where <decode-dir> is where you have the lattices, and is assumed to be"
   echo " a sub-directory of the directory where the model is."
   echo "e.g.: steps/make_index.sh data/kws data/lang exp/sgmm2_5a_mmi/decode/ exp/sgmm2_5a_mmi/decode/kws/"
   echo ""
   echo "main options (for others, see top of script file)"
   echo "  --acwt <float>                                   # acoustic scale used for lattice"
   echo "  --cmd (utils/run.pl|utils/queue.pl <queue opts>) # how to run jobs."
   echo "  --lmwt <float>                                   # lm scale used for lattice"
   echo "  --model <model>                                  # which model to use"
   echo "                                                   # speaker-adapted decoding"
   echo "  --max-silence-frames <int>                       # maximum #frames for silence"
   exit 1;
fi


kwsdatadir=$1;
langdir=$2;
decodedir=$3;
kwsdir=$4;
srcdir=`dirname $decodedir`; # The model directory is one level up from decoding directory.

mkdir -p $kwsdir/log;
nj=`cat $decodedir/num_jobs` || exit 1;
echo $nj > $kwsdir/num_jobs;

utter_id=$kwsdatadir/utter_id
if [ ! -f $utter_id ] ; then
  utter_id=$kwsdatadir/utt.map
fi


if [ -z "$model" ]; then # if --model <mdl> was not specified on the command line...
  model=$srcdir/final.mdl;
fi

for f in $model $decodedir/lat.1.gz $utter_id; do
  [ ! -f $f ] && echo "$0: Error: no such file $f" && exit 1;
done

echo "$0: Using model: $model"

if [ ! -z $silence_word ]; then
  silence_int=`grep -w $silence_word $langdir/words.txt | awk '{print $2}'`
  [ -z $silence_int ] && \
    echo "$0: Error: could not find integer representation of silence word $silence_word" && exit 1;
  silence_opt="--silence-label=$silence_int"
fi

if [ -z "$frame_subsampling_factor" ]; then
  if [ -f $decodedir/../frame_subsampling_factor ] ; then
    frame_subsampling_factor=$(cat $decodedir/../frame_subsampling_factor)
  else 
    frame_subsampling_factor=1
  fi
  echo "$0: Frame subsampling factor autodetected: $frame_subsampling_factor"
fi

word_boundary=$langdir/phones/word_boundary.int
align_lexicon=$langdir/phones/align_lexicon.int
if [ -f $word_boundary ] ; then
  $cmd JOB=1:$nj $kwsdir/log/index.JOB.log \
    lattice-add-penalty --word-ins-penalty=$word_ins_penalty "ark:gzip -cdf $decodedir/lat.JOB.gz|" ark:- \| \
      lattice-align-words $silence_opt --max-expand=$max_expand $word_boundary $model  ark:- ark:- \| \
      lattice-scale --acoustic-scale=$acwt --lm-scale=$lmwt ark:- ark:- \| \
      lattice-to-kws-index --max-states-scale=$max_states_scale --allow-partial=true \
      --frame-subsampling-factor=$frame_subsampling_factor \
      --max-silence-frames=$max_silence_frames --strict=$strict ark:$utter_id ark:- ark:- \| \
      kws-index-union --skip-optimization=$skip_optimization --strict=$strict --max-states=$max_states \
      ark:- "ark:|gzip -c > $kwsdir/index.JOB.gz" || exit 1
elif [ -f $align_lexicon ]; then
  $cmd JOB=1:$nj $kwsdir/log/index.JOB.log \
    lattice-add-penalty --word-ins-penalty=$word_ins_penalty "ark:gzip -cdf $decodedir/lat.JOB.gz|" ark:- \| \
      lattice-align-words-lexicon $silence_opt --max-expand=$max_expand $align_lexicon $model  ark:- ark:- \| \
      lattice-scale --acoustic-scale=$acwt --lm-scale=$lmwt ark:- ark:- \| \
      lattice-to-kws-index --max-states-scale=$max_states_scale --allow-partial=true \
      --frame-subsampling-factor=$frame_subsampling_factor \
      --max-silence-frames=$max_silence_frames --strict=$strict ark:$utter_id ark:- ark:- \| \
      kws-index-union --skip-optimization=$skip_optimization --strict=$strict --max-states=$max_states \
      ark:- "ark:|gzip -c > $kwsdir/index.JOB.gz" || exit 1
else
  echo "$0: Error: cannot find either word-boundary file $word_boundary or alignment lexicon $align_lexicon"
  exit 1
fi

exit 0;
