#!/bin/bash
  # Copyright 2012-2016     Johns Hopkins University (Author: Daniel Povey)
  #                2016     Vimal Manohar
  # Apache 2.0
  
  
  # This script creates biased decoding graphs based on the data transcripts as
  # HCLG.fsts.scp, in the specified directory; this can be consumed by
  # decode_segmentation.sh.
  # This is for use in data-cleanup and data-filtering.
  
  
  set -u
  set -o pipefail
  set -e
  
  # Begin configuration section.
  nj=10
  cmd=run.pl
  scale_opts="--transition-scale=1.0 --self-loop-scale=0.1"
  top_n_words=100 # Number of common words that we compile into each graph (most frequent
                  # in $data/text.orig.
  top_n_words_weight=1.0  # this weight is before renormalization; it can be more
                          # or less than 1.
  min_words_per_graph=100  # Utterances will be grouped so that they have at least
                           # this many words, before making the graph.
  stage=0
  
  ### options for make_one_biased_lm.py.
  ngram_order=4  # maximum n-gram order to use (but see also --min-lm-state-cout).
  min_lm_state_count=10  # make this smaller (e.g. 2) for more strongly biased LM.
  discounting_constant=0.3  # strictly between 0 and 1.  Make this closer to 0 for
                            # more strongly biased LM.
  
  # End configuration options.
  
  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: $0 <data-dir|text> <lang-dir> <dir> <graph-dir>"
     echo "e.g.:  $0 data/train data/lang exp/tri3_cleanup exp/tri3_cleanup/graphs"
     echo "  This script creates biased decoding graphs per utterance (or possibly"
     echo "  groups of utterances, depending on --min-words-per-graph).  Its output"
     echo "  goes to <dir>/HCLG.fsts.scp, indexed by utterance.  Directory <dir> is"
     echo "  required to be a model or alignment directory, containing 'tree' and 'final.mdl'."
     echo "Main options (for others, see top of script file)"
     echo "  --scale-opts <scale-opts>                 # Options relating to language"
     echo "                                            # model scale; default is "
     echo "                                            # '--transition-scale=1.0 --self-loop-scale=0.1'"
     echo "  --top-n-words <N>                         # Number of most-common-words to add with"
     echo "                                            # unigram probabilities into graph (default: 100)"
     echo "  --top-n-words-weight <float>              # Weight given to top-n-words portion of graph"
     echo "                                            # (before renormalizing); may be any positive"
     echo "                                            # number (default: 1.0)"
     echo "  --min-words-per-graph <N>                 # A constant that controls grouping of utterances"
     echo "                                            # (we make the LMs for groups of utterances)."
     echo "                                            # Default: 100."
     echo "  --ngram-order <N>                         # N-gram order in range [2,7].  Maximum n-gram order "
     echo "                                            # that may be used (but also see --min-lm-state-count)."
     echo "                                            # Default 4"
     echo "  --min-lm-state-count <N>                  # Minimum state count for an LM-state of order >2 to "
     echo "                                            # be completely pruned away [bigrams will always be kept]"
     echo "                                            # Default 10.  Smaller -> more strongly biased LM"
     echo "  --discounting-constant <float>            # Discounting constant for Kneser-Ney, strictly between 0"
     echo "                                            # and 1.  Default 0.3.  Smaller -> more strongly biased LM."
     echo "  --config <config-file>                    # config containing options"
     echo "  --nj <nj>                                 # number of parallel jobs"
     echo "  --cmd (utils/run.pl|utils/queue.pl <queue opts>) # how to run jobs."
     exit 1;
  fi
  
  data_or_text=$1
  lang=$2
  dir=$3
  graph_dir=$4
  
  if [ -d $data_or_text ]; then
    text=$data_or_text/text
  else
    text=$data_or_text
  fi
  
  mkdir -p $graph_dir
  
  for f in $text $lang/oov.int $dir/tree $dir/final.mdl \
      $lang/L_disambig.fst $lang/phones/disambig.int; do
    [ ! -f $f ] && echo "$0: expected file $f to exist" && exit 1;
  done
  
  utils/lang/check_phones_compatible.sh $lang/phones.txt $dir/phones.txt
  cp $lang/phones.txt $graph_dir
  
  oov=`cat $lang/oov.int` || exit 1;
  mkdir -p $graph_dir/log
  
  # create top_words.{int,txt}
  if [ $stage -le 0 ]; then
    export LC_ALL=C
    # the following pipe will be broken due to the 'head'; don't fail.
    set +o pipefail
    utils/sym2int.pl --map-oov $oov -f 2- $lang/words.txt < $text | \
      awk '{for(x=2;x<=NF;x++) print $x;}' | sort | uniq -c | \
       sort -nr | head -n $top_n_words > $graph_dir/word_counts.int
    set -o pipefail
    total_count=$(awk '{x+=$1} END{print x}' < $graph_dir/word_counts.int)
    # print top-n words with their unigram probabilities.
    awk -v tot=$total_count -v weight=$top_n_words_weight '{print $2, ($1*weight)/tot;}' \
       <$graph_dir/word_counts.int >$graph_dir/top_words.int
    utils/int2sym.pl -f 1 $lang/words.txt <$graph_dir/top_words.int >$graph_dir/top_words.txt
  fi
  
  word_disambig_symbol=$(cat $lang/words.txt | grep -w "#0" | awk '{print $2}')
  if [ -z "$word_disambig_symbol" ]; then
    echo "$0: error getting word disambiguation symbol"
    exit 1
  fi
  
  mkdir -p $graph_dir/texts
  split_text=
  for n in `seq $nj`; do
    split_text="$split_text $graph_dir/texts/text.$n"
  done
  
  utils/split_scp.pl $text $split_text
  
  mkdir -p $graph_dir/log $graph_dir/fsts
  
  # Make $dir an absolute pathname
  dir=`perl -e '($dir,$pwd)= @ARGV; if($dir!~m:^/:) { $dir = "$pwd/$dir"; } print $dir; ' $dir ${PWD}`
  
  if [ $stage -le 1 ]; then
    echo "$0: creating utterance-group-specific decoding graphs with biased LMs"
  
    # These options are passed through directly to make_one_biased_lm.py.
    lm_opts="--word-disambig-symbol=$word_disambig_symbol --ngram-order=$ngram_order --min-lm-state-count=$min_lm_state_count --discounting-constant=$discounting_constant --top-words=$graph_dir/top_words.int"
  
    $cmd JOB=1:$nj $graph_dir/log/compile_decoding_graphs.JOB.log \
      utils/sym2int.pl --map-oov $oov -f 2- $lang/words.txt $graph_dir/texts/text.JOB \| \
      steps/cleanup/make_biased_lms.py --min-words-per-graph=$min_words_per_graph \
        --lm-opts="$lm_opts" $graph_dir/fsts/utt2group.JOB \| \
      compile-train-graphs-fsts $scale_opts --read-disambig-syms=$lang/phones/disambig.int \
        $dir/tree $dir/final.mdl $lang/L_disambig.fst ark:- \
      ark,scp:$graph_dir/fsts/HCLG.fsts.JOB.ark,$graph_dir/fsts/HCLG.fsts.JOB.scp || exit 1
  fi
  
  for j in $(seq $nj); do cat $graph_dir/fsts/HCLG.fsts.$j.scp; done > $graph_dir/fsts/HCLG.fsts.per_utt.scp
  for j in $(seq $nj); do cat $graph_dir/fsts/utt2group.$j; done > $graph_dir/fsts/utt2group
  
  
  cp $lang/words.txt $graph_dir/
  cp -r $lang/phones $graph_dir/
  
  # The following command gives us an scp file relative to utterance-id.
  utils/apply_map.pl -f 2 $graph_dir/fsts/HCLG.fsts.per_utt.scp <$graph_dir/fsts/utt2group > $graph_dir/HCLG.fsts.scp
  
  n1=$(cat $text | wc -l)
  n2=$(cat $graph_dir/HCLG.fsts.scp | wc -l)
  
  if [ $[$n1*9] -gt $[$n2*10] ]; then
    echo "$0: too many utterances have no scp, something seems to have gone wrong."
    exit 1
  fi
  
  exit 0;