#!/bin/bash
# Copyright 2012-2013  Johns Hopkins University (Author: Daniel Povey);
#                      Arnab Ghoshal
#                2014  Guoguo Chen
#                2015  Hainan Xu

# USAGE:
#
# local/prepare_lang.sh \
#     data_dir/local/dict \
#     data_dir/local/lang \
#     data_dir/lang_dir \
#     "<unk>"

# INPUT:
#
#    "unknown_word"
#
#    data_dir/
#      local/
#        dict/
#          lexicon.txt or lexiconp.txt
#          silence_phones.txt
#          nonsilence_phones.txt
#          optional_silence.txt
#          extra_questions.txt (optional)

# OUTPUT:
# josh@yoga:~/git/kaldi-mirror/egs/kgz/kyrgyz-model$ tree data_org/
# data_org/
# └── lang
#     ├── L_disambig.fst
#     ├── L.fst
#     ├── oov.int
#     ├── oov.txt
#     ├── phones
#     │   ├── align_lexicon.int
#     │   ├── align_lexicon.txt
#     │   ├── context_indep.csl
#     │   ├── context_indep.int
#     │   ├── context_indep.txt
#     │   ├── disambig.csl
#     │   ├── disambig.int
#     │   ├── disambig.txt
#     │   ├── extra_questions.int
#     │   ├── extra_questions.txt
#     │   ├── nonsilence.csl
#     │   ├── nonsilence.int
#     │   ├── nonsilence.txt
#     │   ├── optional_silence.csl
#     │   ├── optional_silence.int
#     │   ├── optional_silence.txt
#     │   ├── roots.int
#     │   ├── roots.txt
#     │   ├── sets.int
#     │   ├── sets.txt
#     │   ├── silence.csl
#     │   ├── silence.int
#     │   └── silence.txt
#     ├── phones.txt
#     ├── topo
#     └── words.txt



# Begin configuration section.
num_sil_states=5
num_nonsil_states=3
position_dependent_phones=true
# position_dependent_phones is false also when position dependent phones and word_boundary.txt 
# have been generated by another source
reverse=false
share_silence_phones=false  # if true, then share pdfs of different silence 
                            # phones together.
sil_prob=0.5
phone_symbol_table=              # if set, use a specified $lang_dir/phones.txt file.
# end configuration sections

. utils/parse_options.sh 

if [ $# -ne 4 ]; then 
    echo "usage: local/prepare_lang.sh <dict-src-dir> <tmp-dir> <lang-dir>  <oov-dict-entry> "
    echo "e.g.: local/prepare_lang.sh data/local/dict <SPOKEN_NOISE> data/local/lang data/lang"
    echo "<dict-src-dir> should contain the following files:"
    echo " extra_questions.txt  lexicon.txt nonsilence_phones.txt  optional_silence.txt  silence_phones.txt"
    echo "See http://kaldi.sourceforge.net/data_prep.html#data_prep_lang_creating for more info."
    echo "options: "
    echo "     --num-sil-states <number of states>             # default: 5, #states in silence models."
    echo "     --num-nonsil-states <number of states>          # default: 3, #states in non-silence models."
    echo "     --position-dependent-phones (true|false)        # default: true; if true, use _B, _E, _S & _I"
    echo "                                                     # markers on phones to indicate word-internal positions. "
    echo "     --reverse (true|false)                          # reverse lexicon."
    echo "     --share-silence-phones (true|false)             # default: false; if true, share pdfs of "
    echo "                                                     # all non-silence phones. "
    echo "     --sil-prob <probability of silence>             # default: 0.5 [must have 0 <= silprob < 1]"
    echo "     --phone-symbol-table <filename>                 # default: \"\"; if not empty, use the provided "
    echo "                                                     # $lang_dir/phones.txt as phone symbol table. This is useful "
    echo "                                                     # if you use a new dictionary for the existing setup."
    exit 1;
fi

src_dir=$1
tmp_dir=$2
lang_dir=$3
oov_word=$4

phone_dir=$lang_dir/phones

# Creating ./$lang_dir and ./$phone_dir
mkdir -p $lang_dir $tmp_dir $phone_dir

silprob=false
[ -f $src_dir/lexiconp_silprob.txt ] && silprob=true

[ -f path.sh ] && . ./path.sh

! utils/validate_dict_dir.pl $src_dir && \
    echo "*Error validating directory $src_dir*" && exit 1;

if [[ ! -f $src_dir/lexicon.txt ]]; then
    echo "**Creating $src_dir/lexicon.txt from $src_dir/lexiconp.txt"
    perl -ape 's/(\S+\s+)\S+\s+(.+)/$1$2/;' < $src_dir/lexiconp.txt > $src_dir/lexicon.txt || exit 1;
fi
if [[ ! -f $src_dir/lexiconp.txt ]]; then
    echo "**Creating $src_dir/lexiconp.txt from $src_dir/lexicon.txt"
    perl -ape 's/(\S+\s+)(.+)/${1}1.0\t$2/;' < $src_dir/lexicon.txt > $src_dir/lexiconp.txt || exit 1;
fi

if ! utils/validate_dict_dir.pl $src_dir >&/dev/null; then
    utils/validate_dict_dir.pl $src_dir  # show the output.
    echo "Validation failed (second time)"
    exit 1;
fi

# $lang_dir/phones.txt file provided, we will do some sanity check here.
if [[ ! -z $phone_symbol_table ]]; then
    # Checks if we have position dependent phones
    n1=`cat $phone_symbol_table | grep -v -E "^#[0-9]+$" | cut -d' ' -f1 | sort -u | wc -l`
    n2=`cat $phone_symbol_table | grep -v -E "^#[0-9]+$" | cut -d' ' -f1 | sed 's/_[BIES]$//g' | sort -u | wc -l`
    $position_dependent_phones && [ $n1 -eq $n2 ] &&\
    echo "$0: Position dependent phones requested, but not in provided phone symbols" && exit 1;
    ! $position_dependent_phones && [ $n1 -ne $n2 ] &&\
      echo "$0: Position dependent phones not requested, but appear in the provided $lang_dir/phones.txt" && exit 1;
    
    # Checks if the phone sets match.
    cat $src_dir/{,non}silence_phones.txt | awk -v f=$phone_symbol_table '
  BEGIN { while ((getline < f) > 0) { sub(/_[BEIS]$/, "", $1); phones[$1] = 1; }}
  { for (x = 1; x <= NF; ++x) { if (!($x in phones)) {
      print "Phone appears in the lexicon but not in the provided $lang_dir/phones.txt: "$x; exit 1; }}}' || exit 1;
fi

if $position_dependent_phones; then
    # Create $tmp_dir/lexicon.original from $src_dir/lexicon.txt by
    # adding the markers _B, _E, _S, _I depending on word position.
    # In this recipe, these markers apply to silence also.
    # Do this starting from lexiconp.txt only.
    if "$silprob"; then 
        perl -ane '@A=split(" ",$_); $w = shift @A; $p = shift @A; $silword_p = shift @A;
              $wordsil_f = shift @A; $wordnonsil_f = shift @A; @A>0||die;
         if(@A==1) { print "$w $p $silword_p $wordsil_f $wordnonsil_f $A[0]_S\n"; } 
         else { print "$w $p $silword_p $wordsil_f $wordnonsil_f $A[0]_B ";
         for($n=1;$n<@A-1;$n++) { print "$A[$n]_I "; } print "$A[$n]_E\n"; } ' \
             < $src_dir/lexiconp_silprob.txt > $tmp_dir/lexiconp_silprob.txt
        if $reverse; then
            echo "We do not support reverse option and silprob at the same time"
            exit 1
        fi
    else
        perl -ane '@A=split(" ",$_); $w = shift @A; $p = shift @A; @A>0||die;
         if(@A==1) { print "$w $p $A[0]_S\n"; } else { print "$w $p $A[0]_B ";
         for($n=1;$n<@A-1;$n++) { print "$A[$n]_I "; } print "$A[$n]_E\n"; } ' \
             < $src_dir/lexiconp.txt > $tmp_dir/lexiconp.pre_reverse || exit 1;
        if $reverse; then
            echo "reversing lexicon."
            cat $tmp_dir/lexiconp.pre_reverse \
                | awk '{printf "%s %s ",$1, $2;for(i=NF;i>2;i--){printf "%s ",$i;}printf "\n"}' \
                > $tmp_dir/lexiconp.txt
        else
            mv $tmp_dir/lexiconp.pre_reverse $tmp_dir/lexiconp.txt
        fi
    fi
    
    # create $tmp_dir/phone_map.txt
    # this has the format (on each line)
    # <original phone> <version 1 of original phone> <version 2> ...
    # where the versions depend on the position of the phone within a word. 
    # For instance, we'd have:
    # AA AA_B AA_E AA_I AA_S
    # for (B)egin, (E)nd, (I)nternal and (S)ingleton
    # and in the case of silence
    # SIL SIL SIL_B SIL_E SIL_I SIL_S
    # [because SIL on its own is one of the variants; this is for when it doesn't
    #  occur inside a word but as an option in the lexicon.]
    
    # This phone map expands the phone lists into all the word-position-dependent
    # versions of the phone lists.
    
    cat <(for x in `cat $src_dir/silence_phones.txt`; do for y in "" "" "_B" "_E" "_I" "_S"; do echo -n "$x$y "; done; echo; done) \
        <(for x in `cat $src_dir/nonsilence_phones.txt`; do for y in "" "_B" "_E" "_I" "_S"; do echo -n "$x$y "; done; echo; done) \
        > $tmp_dir/phone_map.txt
else
    if "$silprob"; then 
        cp $src_dir/lexiconp_silprob.txt $tmp_dir/lexiconp_silprob.txt
        if $reverse; then
            echo "We do not support reverse option and silprob at the same time"
            exit 1
        fi
    else
        cp $src_dir/lexiconp.txt $tmp_dir/lexiconp.pre_reverse
    fi
    
    if $reverse; then
        echo "reversing lexicon."
        cat $tmp_dir/lexiconp.pre_reverse \
            | awk '{printf "%s %s ",$1, $2;for(i=NF;i>2;i--){printf "%s ",$i;}printf "\n"}' \
            > $tmp_dir/lexiconp.txt
    else
        mv $tmp_dir/lexiconp.pre_reverse $tmp_dir/lexiconp.txt
    fi
    
    cat $src_dir/silence_phones.txt $src_dir/nonsilence_phones.txt | \
        sed 's/ /\n/g' | awk '(NF>0){print}' > $tmp_dir/phones
    paste -d' ' $tmp_dir/phones $tmp_dir/phones > $tmp_dir/phone_map.txt
fi

mkdir -p $phone_dir  # various sets of phones...

# Sets of phones for use in clustering, and making monophone systems.

if $share_silence_phones; then
    # build a roots file that will force all the silence phones to share the
    # same pdf's. [three distinct states, only the transitions will differ.]
    # 'shared'/'not-shared' means, do we share the 3 states of the HMM
    # in the same tree-root?
    # Sharing across models(phones) is achieved by writing several phones
    # into one line of roots.txt (shared/not-shared doesn't affect this).
    # 'not-shared not-split' means we have separate tree roots for the 3 states,
    # but we never split the tree so they remain stumps,
    # so all phones in the line correspond to the same model.
    
    cat $src_dir/silence_phones.txt | awk '{printf("%s ", $0); } END{printf("\n");}' | cat - $src_dir/nonsilence_phones.txt | \
        utils/apply_map.pl $tmp_dir/phone_map.txt > $phone_dir/sets.txt
    cat $phone_dir/sets.txt | \
        awk '{if(NR==1) print "not-shared", "not-split", $0; else print "shared", "split", $0;}' > $phone_dir/roots.txt
else
    # different silence phones will have different GMMs.  [note: here, all "shared split" means
    # is that we may have one GMM for all the states, or we can split on states.  because they're
    # context-independent phones, they don't see the context.]
    cat $src_dir/{,non}silence_phones.txt | utils/apply_map.pl $tmp_dir/phone_map.txt > $phone_dir/sets.txt
    cat $phone_dir/sets.txt | awk '{print "shared", "split", $0;}' > $phone_dir/roots.txt
fi

cat $src_dir/silence_phones.txt | utils/apply_map.pl $tmp_dir/phone_map.txt | \
    awk '{for(n=1;n<=NF;n++) print $n;}' > $phone_dir/silence.txt
cat $src_dir/nonsilence_phones.txt | utils/apply_map.pl $tmp_dir/phone_map.txt | \
    awk '{for(n=1;n<=NF;n++) print $n;}' > $phone_dir/nonsilence.txt
cp $src_dir/optional_silence.txt $phone_dir/optional_silence.txt
cp $phone_dir/silence.txt $phone_dir/context_indep.txt

# if extra_questions.txt is empty, it's OK.
cat $src_dir/extra_questions.txt 2>/dev/null | utils/apply_map.pl $tmp_dir/phone_map.txt \
    >$phone_dir/extra_questions.txt

# Want extra questions about the word-start/word-end stuff. Make it separate for
# silence and non-silence. Probably doesn't matter, as silence will rarely
# be inside a word.
if $position_dependent_phones; then
    for suffix in _B _E _I _S; do
        (for x in `cat $src_dir/nonsilence_phones.txt`; do echo -n "$x$suffix "; done; echo) >>$phone_dir/extra_questions.txt
            done
        for suffix in "" _B _E _I _S; do
            (for x in `cat $src_dir/silence_phones.txt`; do echo -n "$x$suffix "; done; echo) >>$phone_dir/extra_questions.txt
                done
fi

# add disambig symbols to the lexicon in $tmp_dir/lexiconp.txt
# and produce $tmp_dir/lexicon_*disambig.txt

if "$silprob"; then
    ndisambig=`utils/add_lex_disambig.pl --pron-probs --sil-probs $tmp_dir/lexiconp_silprob.txt $tmp_dir/lexiconp_silprob_disambig.txt`
else
    ndisambig=`utils/add_lex_disambig.pl --pron-probs $tmp_dir/lexiconp.txt $tmp_dir/lexiconp_disambig.txt`
fi
ndisambig=$[$ndisambig+1]; # add one disambig symbol for silence in lexicon FST.
echo $ndisambig > $tmp_dir/lex_ndisambig

# Format of lexiconp_disambig.txt:
# !SIL	1.0   SIL_S
# <SPOKEN_NOISE>	1.0   SPN_S #1
# <UNK>	1.0  SPN_S #2
# <NOISE>	1.0  NSN_S
# !EXCLAMATION-POINT	1.0  EH2_B K_I S_I K_I L_I AH0_I M_I EY1_I SH_I AH0_I N_I P_I OY2_I N_I T_E

( for n in `seq 0 $ndisambig`; do echo '#'$n; done ) >$phone_dir/disambig.txt

# Create phone symbol table.
if [[ ! -z $phone_symbol_table ]]; then
    start_symbol=`grep \#0 $phone_symbol_table | awk '{print $2}'`
    echo "<eps>" | cat - $phone_dir/{silence,nonsilence}.txt | awk -v f=$phone_symbol_table '
  BEGIN { while ((getline < f) > 0) { phones[$1] = $2; }} { print $1" "phones[$1]; }' | sort -k2 -g |\
    cat - <(cat $phone_dir/disambig.txt | awk -v x=$start_symbol '{n=x+NR-1; print $1, n;}') > $lang_dir/phones.txt 
else
    echo "<eps>" | cat - $phone_dir/{silence,nonsilence,disambig}.txt | \
        awk '{n=NR-1; print $1, n;}' > $lang_dir/phones.txt
fi

# Create a file that describes the word-boundary information for
# each phone.  5 categories.
if $position_dependent_phones; then
    cat $phone_dir/{silence,nonsilence}.txt | \
        awk '/_I$/{print $1, "internal"; next;} /_B$/{print $1, "begin"; next; }
         /_S$/{print $1, "singleton"; next;} /_E$/{print $1, "end"; next; }
         {print $1, "nonword";} ' > $phone_dir/word_boundary.txt
else
    # word_boundary.txt might have been generated by another source
    [ -f $src_dir/word_boundary.txt ] && cp $src_dir/word_boundary.txt $phone_dir/word_boundary.txt
fi

# Create word symbol table.
# <s> and </s> are only needed due to the need to rescore lattices with
# ConstArpaLm format language model. They do not normally appear in G.fst or
# L.fst.

if "$silprob"; then
    # remove the silprob
    cat $tmp_dir/lexiconp_silprob.txt |\
    awk '{
      for(i=1; i<=NF; i++) {
        if(i!=3 && i!=4 && i!=5) printf("%s\t", $i); if(i==NF) print "";
      }
    }' > $tmp_dir/lexiconp.txt
fi

cat $tmp_dir/lexiconp.txt | awk '{print $1}' | sort | uniq  | awk '
  BEGIN {
    print "<eps> 0";
  } 
  {
    if ($1 == "<s>") {
      print "<s> is in the vocabulary!" | "cat 1>&2"
      exit 1;
    }
    if ($1 == "</s>") {
      print "</s> is in the vocabulary!" | "cat 1>&2"
      exit 1;
    }
    printf("%s %d\n", $1, NR);
  }
  END {
    printf("#0 %d\n", NR+1);
    printf("<s> %d\n", NR+2);
    printf("</s> %d\n", NR+3);
  }' > $lang_dir/words.txt || exit 1;

# format of $lang_dir/words.txt:
#<eps> 0
#!EXCLAMATION-POINT 1
#!SIL 2
#"CLOSE-QUOTE 3
#...

silphone=`cat $src_dir/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;

# create $phone_dir/align_lexicon.{txt,int}.
# This is the new-new style of lexicon aligning.

# First remove pron-probs from the lexicon.
perl -ape 's/(\S+\s+)\S+\s+(.+)/$1$2/;' <$tmp_dir/lexiconp.txt >$tmp_dir/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" >> $tmp_dir/align_lexicon.txt

cat $tmp_dir/align_lexicon.txt | \
    perl -ane '@A = split; print $A[0], " ", join(" ", @A), "\n";' | sort | uniq > $phone_dir/align_lexicon.txt

# create phones/align_lexicon.int
cat $phone_dir/align_lexicon.txt | utils/sym2int.pl -f 3- $lang_dir/phones.txt | \
    utils/sym2int.pl -f 1-2 $lang_dir/words.txt > $phone_dir/align_lexicon.int

# Create the basic L.fst without disambiguation symbols, for use
# in training. 

if $silprob; then
    # Usually it's the same as having a fixed-prob L.fst
    # it matters a little bit in discriminative trainings
    utils/make_lexicon_fst_silprob.pl \
        $tmp_dir/lexiconp_silprob_disambig.txt \
        $src_dir/silprob.txt \
        $silphone '#'$ndisambig | \
        sed 's=\#[0-9][0-9]*=<eps>=g' | \
        fstcompile \
            --isymbols=$lang_dir/phones.txt \
            --osymbols=$lang_dir/words.txt \
            --keep_isymbols=false \
            --keep_osymbols=false |   \
        fstarcsort \
            --sort_type=olabel \
            > $lang_dir/L.fst \
        || exit 1;
else
    utils/make_lexicon_fst.pl \
        --pron-probs $tmp_dir/lexiconp.txt \
        $sil_prob \
        $silphone | \
        fstcompile \
            --isymbols=$lang_dir/phones.txt \
            --osymbols=$lang_dir/words.txt \
            --keep_isymbols=false \
            --keep_osymbols=false | \
        fstarcsort \
            --sort_type=olabel \
            > $lang_dir/L.fst \
        || exit 1;
fi

# The file oov.txt contains a word that we will map any OOVs to during
# training.
echo "$oov_word" > $lang_dir/oov.txt || exit 1;
cat $lang_dir/oov.txt | utils/sym2int.pl $lang_dir/words.txt >$lang_dir/oov.int || exit 1;
# integer version of oov symbol, used in some scripts.

# Create these lists of phones in colon-separated integer list form too, 
# for purposes of being given to programs as command-line options.
for f in silence nonsilence optional_silence disambig context_indep; do
    utils/sym2int.pl $lang_dir/phones.txt <$phone_dir/$f.txt >$phone_dir/$f.int
    utils/sym2int.pl $lang_dir/phones.txt <$phone_dir/$f.txt | \
        awk '{printf(":%d", $1);} END{printf "\n"}' | sed s/:// > $phone_dir/$f.csl || exit 1;
done

for x in sets extra_questions; do
    utils/sym2int.pl $lang_dir/phones.txt <$phone_dir/$x.txt > $phone_dir/$x.int || exit 1;
done

utils/sym2int.pl -f 3- $lang_dir/phones.txt <$phone_dir/roots.txt \
    > $phone_dir/roots.int || exit 1;

#if $position_dependent_phones; then
if [ -f $phone_dir/word_boundary.txt ]; then
    utils/sym2int.pl -f 1 $lang_dir/phones.txt <$phone_dir/word_boundary.txt \
        > $phone_dir/word_boundary.int || exit 1;
fi

# this topo file defines the topology for every HMM, that is, 
# for every phone we specify

silphonelist=`cat $phone_dir/silence.csl`
nonsilphonelist=`cat $phone_dir/nonsilence.csl`
utils/gen_topo.pl \
    $num_nonsil_states \
    $num_sil_states \
    $nonsilphonelist \
    $silphonelist \
    >$lang_dir/topo


# Create the lexicon FST with disambiguation symbols, and put it in lang_decode.
# There is an extra step where we create a loop to "pass through" the
# disambiguation symbols from G.fst.
phone_disambig_symbol=`grep \#0 $lang_dir/phones.txt | awk '{print $2}'`
word_disambig_symbol=`grep \#0 $lang_dir/words.txt | awk '{print $2}'`

if $silprob; then
    utils/make_lexicon_fst_silprob.pl $tmp_dir/lexiconp_silprob_disambig.txt \
        $src_dir/silprob.txt $silphone '#'$ndisambig | fstcompile \
        --isymbols=$lang_dir/phones.txt --osymbols=$lang_dir/words.txt \
        --keep_isymbols=false --keep_osymbols=false | \
        fstaddselfloops \
        "echo $phone_disambig_symbol |" "echo $word_disambig_symbol |" | \
        fstarcsort --sort_type=olabel > $lang_dir/L_disambig.fst || exit 1;
else
    utils/make_lexicon_fst.pl --pron-probs $tmp_dir/lexiconp_disambig.txt \
        $sil_prob $silphone '#'$ndisambig | \
        fstcompile --isymbols=$lang_dir/phones.txt \
        --osymbols=$lang_dir/words.txt --keep_isymbols=false \
        --keep_osymbols=false | \
        fstaddselfloops \
        "echo $phone_disambig_symbol |" "echo $word_disambig_symbol |" | \
        fstarcsort --sort_type=olabel > $lang_dir/L_disambig.fst || exit 1;
fi


# validating output directory
! utils/validate_lang.pl $lang_dir && \
    echo "$(basename $0): error validating output" &&  exit 1;

exit 0;
