#!/bin/bash
  # Copyright 2014  Johns Hopkins University (Author: Yenda Trmal)
  # Copyright 2016  Xiaohui Zhang
  #           2018  Ruizhe Huang
  # Apache 2.0
  
  # This script applies a trained Phonetisarus G2P model to
  # synthesize pronunciations for missing words (i.e., words in
  # transcripts but not the lexicon), and output the expanded lexicon.
  # The user could specify either nbest or pmass option 
  # to determine the number of output pronunciation variants, 
  # or use them together to get the intersection of two options.
  
  # Begin configuration section.  
  stage=0
  nbest=      # Generate up to N, like N=3, pronunciation variants for each word
              # (The maximum size of the nbest list, not considering pruning and taking the prob-mass yet). 
  thresh=5    # Pruning threshold for the n-best list, in (0, 99], which is a -log-probability value.
              # A large threshold makes the nbest list shorter, and less likely to hit the max size.
              # This value corresponds to the weight_threshold in shortest-path.h of openfst.
  pmass=      # Select the top variants from the pruned nbest list,
              # summing up to this total prob-mass for a word.
              # On the "boundary", it's greedy by design, e.g. if pmass = 0.8,
              # and we have prob(pron_1) = 0.5, and prob(pron_2) = 0.4, then we get both.
  # End configuration section.
  
  echo "$0 $@"  # Print the command line for logging
  
  [ -f ./path.sh ] && . ./path.sh; # source the path.
  . utils/parse_options.sh || exit 1;
  
  set -u
  set -e
  
  if [ $# != 3 ]; then
    echo "Usage: $0 [options] <word-list> <g2p-model-dir> <output-dir>"
    echo "... where <word-list> is a list of words whose pronunciation is to be generated."
    echo "          <g2p-model-dir> is a directory used as a target during training of G2P"
    echo "          <output-dir> is the directory where the output lexicon should be stored."
    echo "                       The format of the output lexicon output-dir/lexicon.lex is" 
    echo "                       <word>\t<prob>\t<pronunciation> per line."
    echo "e.g.: $0 --nbest 1 exp/g2p/oov_words.txt exp/g2p exp/g2p/oov_lex"
    echo ""
    echo "main options (for others, see top of script file)"
    echo "  --nbest <int>    # Generate upto N pronunciation variants for each word." 
    echo "  --pmass <float>  # Select the top variants from the pruned nbest list," 
    echo "                   # summing up to this total prob-mass, within [0, 1], for a word." 
    echo "  --thresh <int>   # Pruning threshold for n-best."
    exit 1;
  fi
  
  wordlist=$1
  modeldir=$2
  outdir=$3
  
  model=$modeldir/model.fst
  output_lex=$outdir/lexicon.lex
  mkdir -p $outdir
  
  [ ! -f ${model:-} ] && echo "$0: File $model not found in the directory $modeldir." && exit 1
  [ ! -f $wordlist ] && echo "$0: File $wordlist not found!" && exit 1
  [ -z $pmass ] && [ -z $nbest ] && echo "$0: nbest or/and pmass should be specified." && exit 1;
  if ! phonetisaurus=`which phonetisaurus-apply` ; then
    echo "Phonetisarus was not found !"
    echo "Go to $KALDI_ROOT/tools and execute extras/install_phonetisaurus.sh"
    exit 1
  fi
  
  cp $wordlist $outdir/wordlist.txt
  
  # three options: 1) nbest, 2) pmass, 3) nbest+pmass,
  nbest=${nbest:-20}   # if nbest is not specified, set it to 20, due to Phonetisaurus mechanism
  pmass=${pmass:-1.0}  # if pmass is not specified, set it to 1.0, due to Phonetisaurus mechanism
  
  [[ ! $nbest =~ ^[1-9][0-9]*$ ]] && echo "$0: nbest should be a positive integer." && exit 1;
  
  echo "Applying the G2P model to wordlist $wordlist"
  phonetisaurus-apply --pmass $pmass --nbest $nbest --thresh $thresh \
    --word_list $wordlist --model $model \
    --accumulate --verbose --prob \
    1>$output_lex
  
  echo "Completed. Synthesized lexicon for new words is in $output_lex"
  
  # Some words might have been removed or skipped during the process,
  # let's check it and warn the user if so...
  nlex=`cut -f 1 $output_lex | sort -u | wc -l`
  nwlist=`cut -f 1 $wordlist | sort -u | wc -l`
  if [ $nlex -ne $nwlist ] ; then
    failed_wordlist=$outdir/lexicon.failed
    echo "WARNING: Unable to generate pronunciation for all words. ";
    echo "WARINNG:   Wordlist: $nwlist words"
    echo "WARNING:   Lexicon : $nlex words"
    comm -13 <(cut -f 1 $output_lex | sort -u ) \
             <(cut -f 1 $wordlist | sort -u ) \
             >$failed_wordlist && echo "WARNING: The list of failed words is in $failed_wordlist"
  fi
  exit 0