#!/usr/bin/env bash

# This script demonstrates how to use the grammar-decoding framework to build
# graphs made out of more than one part.  (This version uses word-specific
# silence probabilities). It demonstrates using `fstequivalent`
# that the graph constructed this way is equivalent to what you would create if
# you had the LM all as a single piece.  This uses the command line tools to
# expand to a regular FST (--write-as-grammar=false) In practice you might not
# want do to that, since the result might be large, and since writing the entire
# thing might take too much time.  The code itself allows you to construct these
# GrammarFst objects in lightweight way and decode using them.

# Unfortunately the filenames here are not very well through through.  I hope to
# rework this when I have time.

stage=0
run_g2p=false  # set this to true to run the g2p stuff, it's slow so
               # by default we fake it by providing what it previously output
set -e

. ./path.sh
. utils/parse_options.sh


tree_dir=exp/chain/tree_sp
lang_base=data/lang_basevocab
lang_ext=data/lang_extvocab

# For the purposes of this script we just need a biphone tree and associated
# transition-model for testing, because we're testing it at the graph level,
# i.e. testing equivalence of compiled HCLG graphs; there is no decoding
# involved here.

# We're doing this with the "no-silprobs" dictionary dir for now, as we
# need to write some scripts to support silprobs with this.

# For reference, here is how we could create the 'lang' dir for the
# baseline.
#utils/prepare_lang.sh data/local/dict \
#   "<UNK>" data/local/lang_tmp data/lang

if [ $stage -le 0 ]; then
  cp -r data/local/dict data/local/dict_basevocab
  echo "#nonterm:unk" > data/local/dict_basevocab/nonterminals.txt

  utils/prepare_lang.sh data/local/dict_basevocab \
       "<UNK>" data/local/lang_tmp $lang_base
fi

if [ $stage -le 1 ]; then
  # note: <UNK> does appear in that arpa file, with a reasonable probability
  # (0.0)...  presumably because the vocab that the arpa file was built with was
  # not vast, so there were plenty of OOVs.  It would be possible to adjust its
  # probability with adjust_unk_arpa.pl, but for now we just leave it as-is.
  # The <UNK> appears quite a few times in the ARPA.  In the language model we
  # replaced it with #nonterm:unk, which will later expand to our custom graph
  # of new words.

  # We don't want the #nonterm:unk on the output side of G.fst, or it would
  # appear in the decoded output, so we remove it using the 'fstrmsymbols' command.

  nonterm_unk=$(grep '#nonterm:unk' $lang_base/words.txt | awk '{print $2}')

  gunzip -c  data/local/lm/lm_tgsmall.arpa.gz | \
    sed 's/<UNK>/#nonterm:unk/g' | \
    arpa2fst --disambig-symbol=#0 \
             --read-symbol-table=$lang_base/words.txt - | \
    fstrmsymbols --remove-from-output=true "echo $nonterm_unk|" - $lang_base/G.fst
fi


if [ $stage -le 2 ]; then
  # make the top-level part of the graph.
  utils/mkgraph.sh --self-loop-scale 1.0 $lang_base $tree_dir $tree_dir/extvocab_top
fi

if [ $stage -le 3 ] && $run_g2p; then
  # you may have to do some stuff manually to install sequitur, to get this to work.
  dict=data/local/dict_basevocab
  steps/dict/train_g2p.sh --silence-phones $dict/silence_phones.txt $dict/lexicon.txt  $tree_dir/extvocab_g2p
fi


if [ $stage -le 4 ]; then
  # Create data/local/dict_newvocab as a dict-dir containing just the
  # newly created vocabulary entries (but the same phone list as our old setup, not
  # that it matters)

  mkdir -p $tree_dir/extvocab_lexicon

  # First find a list of words in the test set that are out of vocabulary.
  # Of course this is totally cheating.
  awk -v w=data/lang/words.txt 'BEGIN{while(getline <w) seen[$1] = $1} {for(n=2;n<=NF;n++) if(!($n in seen)) oov[$n] = 1}
                                END{ for(k in oov) print k;}' < data/dev_clean_2/text > $tree_dir/extvocab_lexicon/words
  echo "$0: generating g2p entries for $(wc -l <$tree_dir/extvocab_lexicon/words) words"

  if $run_g2p; then
    steps/dict/apply_g2p.sh $tree_dir/extvocab_lexicon/words $tree_dir/extvocab_g2p  $tree_dir/extvocab_lexicon
  else
    cat <<EOF >$tree_dir/extvocab_lexicon//lexicon.lex
HARDWIGG	0.962436	HH AA1 R D W IH1 G
SUDVESTR	0.162048	S AH1 D V EY1 S T R
SUDVESTR	0.133349	S AH1 D V EH1 S T R
SUDVESTR	0.114376	S AH1 D V EH1 S T ER0
VINOS	0.558345	V IY1 N OW0 Z
VINOS	0.068883	V AY1 N OW0 Z
VINOS	0.068431	V IY1 N OW0 S
DOMA	0.645714	D OW1 M AH0
DOMA	0.118255	D UW1 M AH0
DOMA	0.080682	D OW0 M AH0
GWYNPLAINE'S	0.983053	G W IH1 N P L EY1 N Z
SHIMERDA	0.610922	SH IH0 M EH1 R D AH0
SHIMERDA	0.175678	SH IY0 M EH1 R D AH0
SHIMERDA	0.069785	SH AY1 M ER1 D AH0
MYRDALS	0.479183	M IH1 R D AH0 L Z
MYRDALS	0.135225	M ER1 D AH0 L Z
MYRDALS	0.115478	M IH1 R D L Z
HEUCHERA	0.650042	HH OY1 K IH1 R AH0
HEUCHERA	0.119363	HH OY1 K EH1 R AH0
HEUCHERA	0.077907	HH OY1 K ER0 AH0
IMPARA	0.906222	IH0 M P AA1 R AH0
VERLOC'S	0.564847	V ER0 L AA1 K S
VERLOC'S	0.173540	V ER1 L AH0 K S
VERLOC'S	0.050543	V ER1 L AA1 K S
UNTRUSSING	0.998019	AH0 N T R AH1 S IH0 NG
DARFHULVA	0.317057	D AA2 F UH1 L V AH0
DARFHULVA	0.262882	D AA2 F HH UH1 L V AH0
DARFHULVA	0.064055	D AA2 F HH UW1 L V AH0
FINNACTA	0.594586	F IH1 N AH0 K T AH0
FINNACTA	0.232454	F IH1 N AE1 K T AH0
FINNACTA	0.044733	F IH1 N IH0 K T AH0
YOKUL	0.845279	Y OW1 K AH0 L
YOKUL	0.051082	Y OW2 K AH0 L
YOKUL	0.029435	Y OW0 K AH0 L
CONGAL	0.504228	K AA1 NG G AH0 L
CONGAL	0.151648	K AA2 NG G AH0 L
CONGAL	0.137837	K AH0 N JH AH0 L
DELECTASTI	0.632180	D IH0 L EH0 K T EY1 S T IY0
DELECTASTI	0.203808	D IH0 L EH1 K T EY1 S T IY0
DELECTASTI	0.066722	D IH0 L EH0 K T AE1 S T IY0
YUNDT	0.975077	Y AH1 N T
QUINCI	0.426115	K W IH1 N S IY0
QUINCI	0.369324	K W IH1 N CH IY0
QUINCI	0.064507	K W IY0 N CH IY0
BIRDIKINS	0.856979	B ER1 D IH0 K AH0 N Z
BIRDIKINS	0.045315	B ER1 D AH0 K AH0 N Z
SNEFFELS	0.928413	S N EH1 F AH0 L Z
FJORDUNGR	0.130629	F Y AO1 R D UW0 NG G R
FJORDUNGR	0.125082	F Y AO1 R D AH0 NG G R
FJORDUNGR	0.111035	F Y AO1 R D UH1 NG R
YULKA	0.540253	Y UW1 L K AH0
YULKA	0.295588	Y AH1 L K AH0
YULKA	0.076631	Y UH1 L K AH0
LACQUEY'S	0.987908	L AE1 K IY0 Z
OSSIPON'S	0.651400	AA1 S AH0 P AA2 N Z
OSSIPON'S	0.118444	AA1 S AH0 P AA0 N Z
OSSIPON'S	0.106377	AA1 S AH0 P AH0 N Z
SAKNUSSEMM	0.060270	S AE1 K N AH1 S EH1 M
SAKNUSSEMM	0.044992	S AE1 K N AH0 S EH1 M
SAKNUSSEMM	0.044084	S AA0 K N AH1 S EH1 M
CONGAL'S	0.618287	K AA1 NG G AH0 L Z
CONGAL'S	0.185952	K AA2 NG G AH0 L Z
CONGAL'S	0.115143	K AH0 N G AH0 L Z
TARRINZEAU	0.159153	T AA1 R IY0 N Z OW1
TARRINZEAU	0.136536	T AA1 R AH0 N Z OW1
TARRINZEAU	0.100924	T EH1 R IY0 N Z OW1
SHIMERDAS	0.230819	SH IH0 M EH1 R D AH0 Z
SHIMERDAS	0.216235	SH IH0 M EH1 R D AH0 S
SHIMERDAS	0.073311	SH AY1 M ER1 D AH0 Z
RUGGEDO'S	0.821285	R UW0 JH EY1 D OW0 Z
RUGGEDO'S	0.166825	R AH1 G AH0 D OW0 Z
CORNCAKES	0.934118	K AO1 R N K EY2 K S
VENDHYA	0.616662	V EH0 N D Y AH0
VENDHYA	0.178349	V EH1 N D Y AH0
VENDHYA	0.160768	V AA1 N D Y AH0
GINGLE	0.919815	G IH1 NG G AH0 L
STUPIRTI	0.422653	S T UW0 P IH1 R T IY0
STUPIRTI	0.126925	S T UW1 P IH0 R T IY0
STUPIRTI	0.078422	S T UW1 P AH0 R T IY0
HERBIVORE	0.950887	HH ER1 B IH0 V AO2 R
BRION'S	0.838326	B R AY1 AH0 N Z
BRION'S	0.140310	B R IY0 AH0 N Z
DELAUNAY'S	0.993259	D EH1 L AO0 N EY0 Z
KHOSALA	0.920908	K OW0 S AA1 L AH0
BRANDD	0.827461	B R AE1 N D
BRANDD	0.085646	B R AE2 N D
GARDAR	0.598675	G AA0 R D AA1 R
GARDAR	0.289831	G AA1 R D AA2 R
GARDAR	0.057983	G AA0 R D AA2 R
MACKLEWAIN	0.570209	M AE1 K AH0 L W EY0 N
MACKLEWAIN	0.101477	M AH0 K AH0 L W EY0 N
MACKLEWAIN	0.067905	M AE1 K AH0 L W EY2 N
LIBANO	0.993297	L IY0 B AA1 N OW0
MOLING	0.782578	M OW1 L IH0 NG
MOLING	0.059362	M OW2 L IH0 NG
MOLING	0.056217	M AA1 L IH0 NG
BENNYDECK'S	0.583859	B EH1 N IY0 D EH0 K S
BENNYDECK'S	0.276699	B EH1 N IH0 D EH0 K S
BENNYDECK'S	0.028343	B EH1 N IH0 D IH0 K S
MACKLEWAIN'S	0.615766	M AE1 K AH0 L W EY0 N Z
MACKLEWAIN'S	0.109585	M AH0 K AH0 L W EY0 N Z
MACKLEWAIN'S	0.039423	M AE1 K AH0 L W AH0 N Z
PRESTY	0.616071	P R EH1 S T IY0
PRESTY	0.288701	P R AH0 S T IY0
BREADHOUSE	0.995874	B R EH1 D HH AW2 S
BUZZER'S	0.992495	B AH1 Z ER0 Z
BHUNDA	0.502439	B UW1 N D AH0
BHUNDA	0.267733	B AH0 N D AH0
BHUNDA	0.193772	B UH1 N D AH0
PINKIES	0.998440	P IH1 NG K IY0 Z
TROKE	0.723320	T R OW1 K
TROKE	0.269707	T R OW2 K
OSSIPON	0.728486	AA1 S AH0 P AA2 N
OSSIPON	0.098752	AA1 S AH0 P AH0 N
OSSIPON	0.033957	AA1 S AH0 P AO0 N
RIVERLIKE	0.991731	R IH1 V ER0 L AY2 K
NICLESS	0.478183	N IH1 K L AH0 S
NICLESS	0.159889	N IH0 K L AH0 S
NICLESS	0.120611	N IH1 K L IH0 S
TRAMPE	0.959184	T R AE1 M P
VERLOC	0.610461	V ER0 L AA1 K
VERLOC	0.128479	V ER1 L AH0 K
VERLOC	0.073687	V ER1 L AA0 K
GANNY	0.991703	G AE1 N IY0
AMBROSCH	0.302906	AE0 M B R OW1 SH
AMBROSCH	0.201163	AE0 M B R AO1 SH
AMBROSCH	0.109274	AE1 M B R AO1 SH
FIBI	0.619154	F IH1 B IY0
FIBI	0.163168	F IY1 B IY0
FIBI	0.083443	F AY1 B IY0
IROLG	0.823123	IH0 R OW1 L G
IROLG	0.053196	IH0 R OW1 L JH
IROLG	0.021038	IH0 R OW1 L JH IY1
BALVASTRO	0.251546	B AA0 L V AA1 S T R OW0
BALVASTRO	0.213351	B AE0 L V AE1 S T R OW0
BALVASTRO	0.133005	B AA0 L V AE1 S T R OW0
BOOLOOROO	0.676757	B UW1 L UW1 R UW0
BOOLOOROO	0.173653	B UW1 L UH2 R UW0
BOOLOOROO	0.086501	B UW1 L UH0 R UW0
EOF
  fi

  # extend_lang.sh needs it to have basename 'lexiconp.txt'.
  mv $tree_dir/extvocab_lexicon/lexicon.lex $tree_dir/extvocab_lexicon/lexiconp.txt

  [ -f data/lang_extvocab/G.fst ] && rm data/lang_extvocab/G.fst
  utils/lang/extend_lang.sh  data/lang_basevocab $tree_dir/extvocab_lexicon/lexiconp.txt  data/lang_extvocab
fi

if [ $stage -le 5 ]; then
  # make the G.fst for the extra words.  Just assign equal probabilities to all of
  # them.  The words will all transition from state 1 to 2.
  cat <<EOF > $lang_ext/G.txt
0    1    #nonterm_begin <eps>
2    3    #nonterm_end <eps>
3
EOF
  lexicon=$tree_dir/extvocab_lexicon/lexiconp.txt
  num_words=$(wc -l <$lexicon)
  cost=$(perl -e "print log($num_words)");
  awk -v cost=$cost '{print 1, 2, $1, $1, cost}' <$lexicon >>$lang_ext/G.txt
  fstcompile --isymbols=$lang_ext/words.txt --osymbols=$lang_ext/words.txt <$lang_ext/G.txt | \
    fstarcsort --sort_type=ilabel >$lang_ext/G.fst
fi

if [ $stage -le 6 ]; then
  # make the part of the graph that will be included.
  # Refer to the 'compile-graph' commands in ./simple_demo.sh for how you'd do
  # this in code.
  utils/mkgraph.sh --self-loop-scale 1.0 $lang_ext $tree_dir $tree_dir/extvocab_part
fi

if [ $stage -le 7 ]; then
  offset=$(grep nonterm_bos $lang_ext/phones.txt | awk '{print $2}')
  nonterm_unk=$(grep nonterm:unk $lang_ext/phones.txt | awk '{print $2}')

  mkdir -p $tree_dir/extvocab_combined
  [ -d $tree_dir/extvocab_combined/phones ] && rm -r $tree_dir/extvocab_combined/phones
  # the decoding script expects words.txt and phones/, copy them from the extvocab_part
  # graph directory where they will have suitable values.
  cp -r $tree_dir/extvocab_part/{words.txt,phones.txt,phones/} $tree_dir/extvocab_combined

  # the following, due to --write-as-grammar=false, compiles it into an FST
  # which can be decoded by our normal decoder.
  make-grammar-fst --write-as-grammar=false --nonterm-phones-offset=$offset $tree_dir/extvocab_top/HCLG.fst \
                   $nonterm_unk $tree_dir/extvocab_part/HCLG.fst  $tree_dir/extvocab_combined/HCLG.fst

  # the following compiles it and writes as GrammarFst.  The size is 176M, vs. 182M for HCLG.fst.
  # In other examples, of course the difference might be more.

  make-grammar-fst --write-as-grammar=true --nonterm-phones-offset=$offset $tree_dir/extvocab_top/HCLG.fst \
                $nonterm_unk $tree_dir/extvocab_part/HCLG.fst  $tree_dir/extvocab_combined/HCLG.gra
fi


if [ $stage -le 8 ]; then
  # OK, now we actually decode the test data.  For reference, the command which was used to
  # decode the test data in the current (at the time of writing) chain TDNN system
  # local/chain/run_tdnn.sh (as figured out by running it from that stage), was:
  # steps/nnet3/decode.sh --acwt 1.0 --post-decode-acwt 10.0 --frames-per-chunk 140 --nj 38 \
  #   --cmd "queue.pl --mem 4G --num-threads 4" --online-ivector-dir exp/nnet3/ivectors_dev_clean_2_hires \
  #   exp/chain/tree_sp/graph_tgsmall data/dev_clean_2_hires exp/chain/tdnn1h_sp/decode_tgsmall_dev_clean_2

  # We just replace the graph with the one in $treedir/extvocab_combined.

  steps/nnet3/decode.sh --acwt 1.0 --post-decode-acwt 10.0 --frames-per-chunk 140 --nj 38 \
    --cmd "queue.pl --mem 4G --num-threads 4" --online-ivector-dir exp/nnet3/ivectors_dev_clean_2_hires \
    exp/chain/tree_sp/extvocab_combined data/dev_clean_2_hires exp/chain/tdnn1h_sp/decode_tgsmall_dev_clean_2_ev_comb

  # s5: grep WER exp/chain/tdnn1h_sp/decode_tgsmall_dev_clean_2_ev_comb/wer_* | utils/best_wer.sh
  # %WER 11.42 [ 2300 / 20138, 227 ins, 275 del, 1798 sub ] exp/chain/tdnn1h_sp/decode_tgsmall_dev_clean_2_ev_comb/wer_12_0.0

  #.. versus the baseline below:
  # s5: grep WER exp/chain/tdnn1h_sp/decode_tgsmall_dev_clean_2/wer_* | utils/best_wer.sh
  # %WER 12.01 [ 2418 / 20138, 244 ins, 307 del, 1867 sub ] exp/chain/tdnn1h_sp/decode_tgsmall_dev_clean_2/wer_13_0.0
fi

if [ $stage -le 9 ]; then
 steps/nnet3/decode_grammar.sh --acwt 1.0 --post-decode-acwt 10.0 --frames-per-chunk 140 --nj 38 \
    --cmd "queue.pl --mem 4G --num-threads 4" --online-ivector-dir exp/nnet3/ivectors_dev_clean_2_hires \
    exp/chain/tree_sp/extvocab_combined data/dev_clean_2_hires exp/chain/tdnn1h_sp/decode_tgsmall_dev_clean_2_ev_comb_gra

 # WER with grammar decoding is exactly the same as decoding from the converted FST.
 # grep WER exp/chain/tdnn1h_sp/decode_tgsmall_dev_clean_2_ev_comb_gra/wer_* | utils/best_wer.sh
 # %WER 11.42 [ 2300 / 20138, 227 ins, 275 del, 1798 sub ] exp/chain/tdnn1h_sp/decode_tgsmall_dev_clean_2_ev_comb_gra/wer_12_0.0
fi