#!/bin/bash
  # Copyright     2018  Johns Hopkins University (Author: Daniel Povey);
  
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #  http://www.apache.org/licenses/LICENSE-2.0
  #
  # THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  # KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
  # WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
  # MERCHANTABLITY OR NON-INFRINGEMENT.
  # See the Apache 2 License for the specific language governing permissions and
  # limitations under the License.
  
  # derived files, that go in data/lang/.
  
  # Begin configuration section.
  sil_prob=0.5
  # end configuration section
  
  echo "$0 $@"  # Print the command line for logging
  
  . utils/parse_options.sh
  
  if [ $# -ne 3 ]; then
    echo "Usage: utils/extend_lang.sh <old-lang-dir> <lexicon> <new-lang-dir>"
    echo "e.g.: utils/extend_lang.sh data/lang data/local/dict_new_words/lexiconp.txt data/lang_new_words"
    echo ""
    echo "This script creates a lang/ directory <new-lang-dir> with L.fst and L_disambig.fst"
    echo "derived from the provided lexicon, but all other information being the same as the old"
    echo "lang/ directory, including the phones.txt and words.txt being compatible (however,"
    echo "words.txt may have new words, and phones.txt may have extra disambiguation symbols"
    echo "if needed).  We do not allow new phones."
    echo ""
    echo "CAUTION: the lexicon generated will only cover the words in the provided lexicon,"
    echo "which might not include all the words in words.txt.  You should make sure your"
    echo "lexicon is a superset of the original lexicon used to generate <old-lang-dir>,"
    echo "if this would be a problem for your scenario."
    echo ""
    echo "The basename of <lexicon> must be either lexicon.txt, lexiconp.txt or lexiconp_silprob.txt."
    echo ""
    echo "Options"
    echo "     --sil-prob <probability of silence>             # default: 0.5 [must have 0 <= silprob < 1]"
    exit 1;
  fi
  
  srcdir=$1
  lexicon=$2
  dir=$3
  
  [ -f path.sh ] && . ./path.sh
  
  for f in $srcdir/phones.txt $lexicon; do
    if [ ! -f $f ]; then
      echo "$0: expected file $f to exist"
      exit 1
    fi
  done
  
  if ! awk '{if(NF < 2) exit(1)} END{if(NR==0) exit(1)}' <$lexicon; then
    echo "$0: it looks like there words without pronunciations or.."
    echo "  ...blank lines in $lexicon, or it is empty."
    exit 1
  fi
  
  mkdir -p $dir
  
  if [ -d $dir/phones ]; then rm -r $dir/phones; fi
  
  cp -r $srcdir/phones $dir/
  
  for f in oov.int oov.txt phones.txt topo words.txt; do
    cp $srcdir/$f $dir/
  done
  
  tmpdir=$dir/temp
  rm -r $tmpdir 2>/dev/null
  mkdir -p $tmpdir
  
  
  # TODO: more checking.
  if [ $(basename $lexicon) != lexiconp.txt ]; then
    echo "$0: currently this script only supports the lexiconp.txt format; your lexicon"
    echo " ... has to have that filename."
  fi
  
  # Get the list of extra words.
  awk -v w=$srcdir/words.txt 'BEGIN{while(getline <w) seen[$1] = $1} { if (!($1 in seen)) oov[$1] = 1}
                       END{ for(k in oov) print k;}' <$lexicon >$tmpdir/extra_words.txt
  
  # Add entries to words.txt for all the words that were not previously in the
  # lexicon.
  highest_number=$(tail -n 1 $srcdir/words.txt | awk '{print $2}')
  awk -v start=$highest_number '{print $1, NR+start}' <$tmpdir/extra_words.txt >>$dir/words.txt
  echo "$0: added $(wc -l <$tmpdir/extra_words.txt) extra words to words.txt"
  
  if [ -f $dir/phones/nonterminals.txt ]; then
    # extra grammar-decoding-related options for getting the lexicon.
    grammar_opts="--left-context-phones=$dir/phones/left_context_phones.txt --nonterminals=$srcdir/phones/nonterminals.txt"
  else
    grammar_opts=""
  fi
  
  if [ -f $dir/phones/word_boundary.txt ]; then
    # was `if $position_dependent_phones; then..` in prepare_lang.sh
    # TODO: add support for silprobs
    perl -ane '@A=split(" ",$_); $w = shift @A; $p = shift @A; @A>0||die;
           if(@A==1) { print "$w $p $A[0]_S
  "; } else { print "$w $p $A[0]_B ";
           for($n=1;$n<@A-1;$n++) { print "$A[$n]_I "; } print "$A[$n]_E
  "; } ' \
         < $lexicon > $tmpdir/lexiconp.txt || exit 1;
  else
    cp $lexicon $tmpdir/lexiconp.txt
  fi
  
  # Check that there are no unseen phones in the lexicon.
  if ! utils/sym2int.pl -f 3- $srcdir/phones.txt $tmpdir/lexiconp.txt >/dev/null; then
    echo "$0: it looks like there are unseen phones in your lexicon $lexicon"
    exit 1
  fi
  
  ndisambig=$(utils/add_lex_disambig.pl --pron-probs $tmpdir/lexiconp.txt $tmpdir/lexiconp_disambig.txt)
  
  ndisambig=$[ndisambig+1]  # Add one to disambiguate silence.
  
  # we'll need to figure out whether any of these disambiguation symbols are
  # absent from our current disambiguation phones.. if they are, then we need to
  # add them as new disambiguation symbols to phones.txt.
  for n in $(seq 0 $ndisambig); do
    sym='#'$n; if ! grep -w -q "$sym" $dir/phones/disambig.txt; then echo "$sym"; fi
  done > $tmpdir/extra_disambig.txt
  highest_number=$(tail -n 1 $srcdir/phones.txt | awk '{print $2}')
  awk -v start=$highest_number '{print $1, NR+start}' <$tmpdir/extra_disambig.txt >>$dir/phones.txt
  echo "$0: added $(wc -l <$tmpdir/extra_disambig.txt) extra disambiguation symbols to phones.txt"
  
  # add extra_disambig symbols into disambig.txt
  cat $tmpdir/extra_disambig.txt >> $dir/phones/disambig.txt
  utils/sym2int.pl $dir/phones.txt <$dir/phones/disambig.txt >$dir/phones/disambig.int
  utils/sym2int.pl $dir/phones.txt <$dir/phones/disambig.txt | \
    awk '{printf(":%d", $1);} END{printf "
  "}' | sed s/:// > $dir/phones/disambig.csl
  
  silphone=`cat $srcdir/phones/optional_silence.txt` || exit 1;
  [ -z "$silphone" ] && \
    ( echo "You have no optional-silence phone; it is required in the current scripts"
      echo "but you may use the option --sil-prob 0.0 to stop it being used." ) && \
     exit 1;
  
  
  # First remove pron-probs from the lexicon.
  perl -ape 's/(\S+\s+)\S+\s+(.+)/$1$2/;' <$tmpdir/lexiconp.txt >$tmpdir/align_lexicon.txt
  
  # Note: here, $silphone will have no suffix e.g. _S because it occurs as optional-silence,
  # and is not part of a word.
  [ ! -z "$silphone" ] && echo "<eps> $silphone" >> $tmpdir/align_lexicon.txt
  
  cat $tmpdir/align_lexicon.txt | \
    perl -ane '@A = split; print $A[0], " ", join(" ", @A), "
  ";' | sort | uniq > $dir/phones/align_lexicon.txt
  
  # create phones/align_lexicon.int from phones/align_lexicon.txt
  cat $dir/phones/align_lexicon.txt | utils/sym2int.pl -f 3- $dir/phones.txt | \
    utils/sym2int.pl -f 1-2 $dir/words.txt > $dir/phones/align_lexicon.int
  
  # Create the basic L.fst without disambiguation symbols, for use
  # in training.
  
  utils/lang/make_lexicon_fst.py $grammar_opts --sil-prob=$sil_prob --sil-phone=$silphone \
           $tmpdir/lexiconp.txt | \
      fstcompile --isymbols=$dir/phones.txt --osymbols=$dir/words.txt \
        --keep_isymbols=false --keep_osymbols=false | \
      fstarcsort --sort_type=olabel > $dir/L.fst || exit 1;
  
  
  # and create the version that has disambiguation symbols.
  utils/lang/make_lexicon_fst.py $grammar_opts \
    --sil-prob=$sil_prob --sil-phone=$silphone --sil-disambig='#'$ndisambig \
       $tmpdir/lexiconp_disambig.txt | \
     fstcompile --isymbols=$dir/phones.txt --osymbols=$dir/words.txt \
                --keep_isymbols=false --keep_osymbols=false | \
     fstaddselfloops $dir/phones/wdisambig_phones.int $dir/phones/wdisambig_words.int | \
     fstarcsort --sort_type=olabel > $dir/L_disambig.fst || exit 1;
  
  
  echo "$(basename $0): validating output directory"
  # the --skip-generate-words-check option is needed because L.fst may not actually
  # contain all the words in words.txt.
  ! utils/validate_lang.pl --skip-generate-words-check $dir && echo "$(basename $0): error validating output" &&  exit 1;
  
  exit 0;