#!/bin/bash
  # Copyright 2012-2018  Johns Hopkins University (Author: Guoguo Chen, Yenda Trmal)
  # License: Apache 2.0
  
  
  help_message="$(basename $0): do keyword indexing and search.  data-dir is assumed to have
                   kws/ subdirectory that specifies the terms to search for.  Output is in
                   decode-dir/kws/
               Usage:
                   $(basename $0) <lang-dir> <data-dir> <decode-dir>"
  
  # Begin configuration section.
  min_lmwt=8
  max_lmwt=12
  cmd=run.pl
  model=
  skip_scoring=false
  skip_optimization=false # true can speed it up if #keywords is small.
  max_states=350000
  indices_dir=
  kwsout_dir=
  stage=0
  word_ins_penalty=0
  extraid=
  silence_word=  # specify this if you did to in kws_setup.sh, it's more accurate.
  strict=false
  duptime=0.6
  ntrue_scale=1.0
  frame_subsampling_factor=1
  nbest=-1
  max_silence_frames=50
  skip_indexing=false
  # End configuration section.
  
  echo "$0 $@"  # Print the command line for logging
  
  [ -f ./path.sh ] && . ./path.sh; # source the path.
  . parse_options.sh || exit 1;
  
  set -u
  set -e
  set -o pipefail
  
  
  if [[ "$#" -ne "3" ]] ; then
      echo -e "$0: FATAL: wrong number of script parameters!
  
  "
      printf "$help_message
  
  "
      exit 1;
  fi
  
  silence_opt=
  
  langdir=$1
  datadir=$2
  decodedir=$3
  
  if [ -z $extraid ] ; then
    kwsdatadir=$datadir/kws
  else
    kwsdatadir=$datadir/kwset_${extraid}
  fi
  
  if [ -z $extraid ] ; then
    kwsoutdir=$decodedir/kws
  else
    kwsoutdir=$decodedir/kwset_${extraid}
  fi
  
  
  if [ -z $indices_dir ]; then
    indices_dir=$kwsoutdir
  fi
  
  if [ ! -z "$model" ]; then
      model_flags="--model $model"
  else
      model_flags=
  fi
  
  mkdir -p $kwsoutdir
  for d in "$datadir" "$kwsdatadir" "$langdir" "$decodedir"; do
    if [ ! -d "$d" ]; then
      echo "$0: FATAL: expected directory $d to exist"
      exit 1;
    fi
  done
  
  echo "$0: Searching: $kwsdatadir"
  duration=$(cat $kwsdatadir/trials)
  echo "$0: Duration: $duration"
  
  
  frame_subsampling_factor=1
  if [ -f $decodedir/../frame_subsampling_factor ] ; then
    frame_subsampling_factor=$(cat $decodedir/../frame_subsampling_factor)
    echo "$0: Frame subsampling factor autodetected: $frame_subsampling_factor"
  elif [ -f $decodedir/../../frame_subsampling_factor ] ; then
    frame_subsampling_factor=$(cat $decodedir/../../frame_subsampling_factor)
    echo "$0: Frame subsampling factor autodetected: $frame_subsampling_factor"
  fi
  
  if [ $stage -le 0 ] ; then
    if [ ! -f $indices_dir/.done.index ] && ! $skip_indexing ; then
      [ ! -d $indices_dir ] && mkdir  $indices_dir
      for lmwt in $(seq $min_lmwt $max_lmwt) ; do
        indices=${indices_dir}_$lmwt
        mkdir -p $indices
  
        acwt=$(perl -e "print 1.0/$lmwt")
        [ ! -z $silence_word ] && silence_opt="--silence-word $silence_word"
        steps/make_index.sh $silence_opt --cmd "$cmd" --acwt $acwt $model_flags\
          --skip-optimization $skip_optimization --max-states $max_states \
          --word-ins-penalty $word_ins_penalty --max-silence-frames $max_silence_frames\
          --frame-subsampling-factor ${frame_subsampling_factor} \
          $kwsdatadir $langdir $decodedir $indices  || exit 1
      done
      touch $indices_dir/.done.index
    else
      echo "$0: Assuming indexing has been aready done. If you really need to re-run "
      echo "$0: the indexing again, delete the file $indices_dir/.done.index"
    fi
  fi
  
  keywords=$kwsdatadir/keywords.fsts
  if [ -f $keywords ] ; then
    echo "$0: Using ${keywords} for search"
    keywords="ark:$keywords"
  elif [ -f ${keywords}.gz ] ; then
    echo "$0: Using ${keywords}.gz for search"
    keywords="ark:gunzip -c ${keywords}.gz |"
  else
    echo "$0: The keyword file ${keywords}[.gz] does not exist"
  fi
  
  
  if [ $stage -le 1 ]; then
    for lmwt in $(seq $min_lmwt $max_lmwt) ; do
      kwsoutput=${kwsoutdir}_$lmwt
      indices=${indices_dir}_$lmwt
      nj=$(cat $indices/num_jobs)
  
  
      for f in $indices/index.1.gz ; do
        [ ! -f $f ] && echo "$0: no such file $f" && exit 1;
      done
  
      mkdir -p $kwsoutput/log
      $cmd JOB=1:$nj $kwsoutput/log/search.JOB.log \
        set -e  -o pipefail '&&' \
        kws-search --strict=$strict --negative-tolerance=-1 \
        --frame-subsampling-factor=${frame_subsampling_factor} \
        "ark:gzip -cdf $indices/index.JOB.gz|" "$keywords" \
        "ark,t:| sort -u | gzip -c > $kwsoutput/result.JOB.gz" \
        "ark,t:| sort -u | gzip -c > $kwsoutput/stats.JOB.gz" || exit 1;
    done
  fi
  
  if [ $stage -le 2 ]; then
    for lmwt in $(seq $min_lmwt $max_lmwt) ; do
      kwsoutput=${kwsoutdir}_$lmwt
      indices=${indices_dir}_$lmwt
      nj=$(cat $indices/num_jobs)
  
      # This is a memory-efficient way how to do the filtration
      # we do this in this way because the result.* files can be fairly big
      # and we do not want to run into troubles with memory
      files=""
      for job in $(seq 1 $nj); do
        if [ -f $kwsoutput/result.${job}.gz ] ; then
         files="$files <(gunzip -c $kwsoutput/result.${job}.gz)"
        elif [ -f $kwsoutput/result.${job} ] ; then
         files="$files $kwsoutput/result.${job}"
        else
          echo >&2 "The file $kwsoutput/result.${job}[.gz] does not exist"
          exit 1
        fi
      done
      # we have to call it using eval as we need the bash to interpret
      # the (possible) command substitution in case of gz files
      # bash -c would probably work as well, but would spawn another
      # shell instance
      eval "sort -m -u $files" |\
        local/kws/filter_kws_results.pl --likes --nbest $nbest > $kwsoutput/results || exit 1
    done
  fi
  
  if [ -z $extraid ] ; then
    extraid_flags=
  else
    extraid_flags="  --extraid ""$extraid"" "
  fi
  
  if [ $stage -le 4 ]; then
    if $skip_scoring ; then
      echo "$0: Not scoring, because --skip-scoring true was issued"
    elif [ ! -x local/kws/score.sh ] ; then
      echo "$0: Not scoring, because the file local/kws_score.sh is not present"
    else
      echo "$0: Scoring KWS results"
      local/kws/score.sh --cmd "$cmd" \
        --min-lmwt $min_lmwt --max-lmwt $max_lmwt $extraid_flags \
        $langdir $datadir ${kwsoutdir} || exit 1;
    fi
  fi
  
  echo "$0: Done"
  exit 0