#!/bin/bash

# Copyright 2016  Vimal Manohar
# Apache 2.0

# This script demonstrates how to re-segment long audios into short segments.
# The basic idea is to decode with an existing in-domain acoustic model, and a
# bigram language model built from the reference, and then work out the
# segmentation from a ctm like file.

## %WER results. 

## Baseline results
# %WER 18.1 | 507 17783 | 84.7 10.7 4.6 2.8 18.1 91.1 | -0.073 | exp/tri3/decode_dev_rescore/score_16_0.0/ctm.filt.filt.sys
# %WER 16.6 | 1155 27500 | 85.7 10.7 3.6 2.4 16.6 86.0 | -0.041 | exp/tri3/decode_test_rescore/score_16_0.0/ctm.filt.filt.sys

## With Cleanup
# %WER 18.0 | 507 17783 | 85.0 10.6 4.4 3.0 18.0 90.9 | -0.064 | exp/tri3_cleaned/decode_dev_rescore/score_14_0.0/ctm.filt.filt.sys
# %WER 16.6 | 1155 27500 | 85.9 10.8 3.3 2.5 16.6 86.6 | -0.050 | exp/tri3_cleaned/decode_test_rescore/score_14_0.0/ctm.filt.filt.sys

## Segmentation results
# %WER 18.9 | 507 17783 | 83.9 11.1 5.0 2.8 18.9 92.9 | -0.103 | exp/tri3_reseg_a/decode_nosp_dev_rescore/score_14_0.0/ctm.filt.filt.sys
# %WER 17.6 | 1155 27500 | 84.6 11.3 4.1 2.2 17.6 86.8 | -0.005 | exp/tri3_reseg_a/decode_nosp_test_rescore/score_14_0.0/ctm.filt.filt.sys

## Segmentation + Cleanup

# cleaned - 
# Default segmentation-opts "--max-junk-proportion=1 --max-deleted-words-kept-when-merging=1 --min-split-point-duration=0.1" 
# cleaned_b -
# "--max-junk-proportion=0.5 --max-deleted-words-kept-when-merging=10"
# cleaned_c -
# "--max-junk-proportion=0.2 --max-deleted-words-kept-when-merging=6 --min-split-point-duration=0.3"

# %WER 18.7 | 507 17783 | 84.0 11.0 5.0 2.8 18.7 91.7 | -0.119 | exp/tri3_reseg_a_cleaned/decode_nosp_dev_rescore/score_15_0.0/ctm.filt.filt.sys
# %WER 18.6 | 507 17783 | 84.0 11.0 4.9 2.7 18.6 91.5 | -0.092 | exp/tri3_reseg_a_cleaned_b/decode_nosp_dev_rescore/score_15_0.0/ctm.filt.filt.sys
# %WER 18.6 | 507 17783 | 84.1 10.8 5.0 2.7 18.6 92.1 | -0.114 | exp/tri3_reseg_a_cleaned_c/decode_nosp_dev_rescore/score_15_0.0/ctm.filt.filt.sys

# %WER 17.7 | 1155 27500 | 84.5 11.4 4.0 2.2 17.7 86.8 | -0.020 | exp/tri3_reseg_a_cleaned/decode_nosp_test_rescore/score_14_0.0/ctm.filt.filt.sys
# %WER 17.3 | 1155 27500 | 84.8 11.2 4.1 2.1 17.3 86.8 | -0.002 | exp/tri3_reseg_a_cleaned_b/decode_nosp_test_rescore/score_15_0.0/ctm.filt.filt.sys
# %WER 17.7 | 1155 27500 | 84.6 11.4 4.1 2.3 17.7 86.6 | -0.018 | exp/tri3_reseg_a_cleaned_c/decode_nosp_test_rescore/score_14_0.0/ctm.filt.filt.sys

## Use silence and pronunciation probs estimated from resegmented data
# %WER 18.2 | 507 17783 | 84.6 10.8 4.5 2.9 18.2 92.5 | -0.037 | exp/tri3_reseg_a/decode_a_dev_rescore/score_16_0.0/ctm.filt.filt.sys
# %WER 16.9 | 1155 27500 | 85.5 11.0 3.5 2.4 16.9 86.1 | -0.024 | exp/tri3_reseg_a/decode_a_test_rescore/score_14_0.0/ctm.filt.filt.sys

## Use silence and pronunciation probs estimated from resegmented and cleaned up data
# %WER 18.2 | 507 17783 | 84.4 10.8 4.9 2.6 18.2 92.5 | -0.074 | exp/tri3_reseg_a_cleaned_b/decode_a_cleaned_b_dev_rescore/score_15_0.5/ctm.filt.filt.sys
# %WER 16.8 | 1155 27500 | 85.4 10.7 3.9 2.1 16.8 86.8 | -0.046 | exp/tri3_reseg_a_cleaned_b/decode_a_cleaned_b_test_rescore/score_14_0.5/ctm.filt.filt.sys

. ./cmd.sh
. ./path.sh

set -e -o pipefail -u

segment_stage=-9
cleanup_stage=-1
cleanup_affix=cleaned_b
affix=_a

decode_nj=8   # note: should not be >38 which is the number of speakers in the dev set
              # after applying --seconds-per-spk-max 180.  We decode with 4 threads, so
              # this will be too many jobs if you're using run.pl.

###############################################################################
# Simulate unsegmented data directory.
###############################################################################
utils/data/convert_data_dir_to_whole.sh data/train data/train_long

###############################################################################
# Train system on a small subset of 2000 utterances that are 
# manually segmented.
###############################################################################

utils/subset_data_dir.sh --speakers data/train 2000 data/train_2k
utils/subset_data_dir.sh --shortest data/train_2k 500 data/train_2k_short500

steps/make_mfcc.sh --cmd "$train_cmd" --nj 32 \
  data/train_long exp/make_mfcc/train_long mfcc || exit 1
steps/compute_cmvn_stats.sh data/train_long \
  exp/make_mfcc/train_long mfcc


steps/train_mono.sh --nj 20 --cmd "$train_cmd" \
  data/train_2k_short500 data/lang_nosp exp/mono_a

steps/align_si.sh --nj 20 --cmd "$train_cmd" \
  data/train_2k data/lang_nosp exp/mono_a exp/mono_a_ali_2k

steps/train_deltas.sh --cmd "$train_cmd" \
  500 5000 data/train_2k data/lang_nosp exp/mono_a_ali_2k exp/tri1a

steps/align_si.sh --nj 20 --cmd "$train_cmd" \
  data/train_2k data/lang_nosp exp/tri1a exp/tri1a_ali

steps/train_lda_mllt.sh --cmd "$train_cmd" \
  1000 10000 data/train_2k data/lang_nosp exp/tri1a_ali exp/tri1b

###############################################################################
# Segment long recordings using TF-IDF retrieval of reference text 
# for uniformly segmented audio chunks based on Smith-Waterman alignment.
# Use a model trained on train_2k (tri1b)
###############################################################################

steps/cleanup/segment_long_utterances.sh --cmd "$train_cmd" \
  --stage $segment_stage --nj 80 \
  --max-bad-proportion 0.5 \
  exp/tri1b data/lang_nosp data/train_long data/train_reseg${affix} \
  exp/segment_long_utts${affix}_train

steps/compute_cmvn_stats.sh data/train_reseg${affix} \
  exp/make_mfcc/train_reseg${affix} mfcc
utils/fix_data_dir.sh data/train_reseg${affix}

###############################################################################
# Train new model on segmented data directory starting from the same model
# used for segmentation. (tri2_reseg)
###############################################################################

# Align tri1b system with reseg${affix} data
steps/align_si.sh  --nj 40 --cmd "$train_cmd" \
  data/train_reseg${affix} \
  data/lang_nosp exp/tri1b exp/tri1b_ali_reseg${affix}  || exit 1;

# Train LDA+MLLT system on reseg${affix} data
steps/train_lda_mllt.sh --cmd "$train_cmd" \
  4000 50000 data/train_reseg${affix} data/lang_nosp \
  exp/tri1b_ali_reseg${affix} exp/tri2_reseg${affix}

(
utils/mkgraph.sh data/lang_nosp exp/tri2_reseg${affix} \
  exp/tri2_reseg${affix}/graph_nosp
for dset in dev test; do
  steps/decode.sh --nj $decode_nj --cmd "$decode_cmd"  --num-threads 4 \
    exp/tri2_reseg${affix}/graph_nosp data/${dset} \
    exp/tri2_reseg${affix}/decode_nosp_${dset}
  steps/lmrescore_const_arpa.sh  --cmd "$decode_cmd" data/lang_nosp \
    data/lang_nosp_rescore \
     data/${dset} exp/tri2_reseg${affix}/decode_nosp_${dset} \
     exp/tri2_reseg${affix}/decode_nosp_${dset}_rescore
done
) &

###############################################################################
# Train SAT model on segmented data directory
###############################################################################

# Train SAT system on reseg${affix} data
steps/train_sat.sh --cmd "$train_cmd" 5000 100000 \
  data/train_reseg${affix} data/lang_nosp \
  exp/tri2_reseg${affix} exp/tri3_reseg${affix}

(
utils/mkgraph.sh data/lang_nosp exp/tri3_reseg${affix} \
  exp/tri3_reseg${affix}/graph_nosp
for dset in dev test; do
  steps/decode_fmllr.sh --nj $decode_nj --cmd "$decode_cmd"  --num-threads 4 \
    exp/tri3_reseg${affix}/graph_nosp data/${dset} \
    exp/tri3_reseg${affix}/decode_nosp_${dset}
  steps/lmrescore_const_arpa.sh  --cmd "$decode_cmd" data/lang_nosp \
    data/lang_nosp_rescore \
     data/${dset} exp/tri3_reseg${affix}/decode_nosp_${dset} \
     exp/tri3_reseg${affix}/decode_nosp_${dset}_rescore
done
) &

###############################################################################
# Clean and segmented data
###############################################################################

segmentation_opts=(
--max-junk-proportion=0.5
--max-deleted-words-kept-when-merging=10
)
opts="${segmentation_opts[@]}"

steps/cleanup/clean_and_segment_data.sh --nj 40 --cmd "$train_cmd" \
  --segmentation-opts "$opts" \
  data/train_reseg${affix} data/lang_nosp exp/tri3_reseg${affix} \
  exp/tri3_reseg${affix}_${cleanup_affix}_work \
  data/train_reseg${affix}_${cleanup_affix}

###############################################################################
# Train new SAT model on cleaned data directory 
###############################################################################

steps/align_fmllr.sh --nj 40 --cmd "$train_cmd" \
  data/train_reseg${affix}_${cleanup_affix} data/lang_nosp \
  exp/tri3_reseg${affix} exp/tri3_reseg${affix}_ali_${cleanup_affix}

steps/train_sat.sh --cmd "$train_cmd" \
  5000 100000 data/train_reseg${affix}_${cleanup_affix} data/lang_nosp \
  exp/tri3_reseg${affix}_ali_${cleanup_affix} \
  exp/tri3_reseg${affix}_$cleanup_affix

(
utils/mkgraph.sh data/lang_nosp exp/tri3_reseg${affix}_$cleanup_affix \
  exp/tri3_reseg${affix}_$cleanup_affix/graph_nosp
for dset in dev test; do
  steps/decode_fmllr.sh --nj $decode_nj --cmd "$decode_cmd"  --num-threads 4 \
    exp/tri3_reseg${affix}_$cleanup_affix/graph_nosp data/${dset} \
    exp/tri3_reseg${affix}_$cleanup_affix/decode_nosp_${dset}
  steps/lmrescore_const_arpa.sh  --cmd "$decode_cmd" data/lang_nosp \
    data/lang_nosp_rescore \
     data/${dset} exp/tri3_reseg${affix}_$cleanup_affix/decode_nosp_${dset} \
     exp/tri3_reseg${affix}_$cleanup_affix/decode_nosp_${dset}_rescore
done
) &

steps/get_prons.sh --cmd "$train_cmd" \
  data/train_reseg${affix}_${cleanup_affix} \
  data/lang_nosp exp/tri3_reseg${affix}_$cleanup_affix
utils/dict_dir_add_pronprobs.sh --max-normalize true \
  data/local/dict_nosp \
  exp/tri3_reseg${affix}_$cleanup_affix/{pron,sil,pron_bigram}_counts_nowb.txt \
  data/local/dict${affix}_$cleanup_affix

utils/prepare_lang.sh data/local/dict${affix}_$cleanup_affix \
  "<unk>" data/local/lang data/lang${affix}_$cleanup_affix
cp -rT data/lang${affix}_$cleanup_affix data/lang${affix}_${cleanup_affix}_rescore
cp data/lang_nosp/G.fst data/lang${affix}_$cleanup_affix/
cp data/lang_nosp_rescore/G.carpa data/lang${affix}_${cleanup_affix}_rescore/

(
utils/mkgraph.sh data/lang${affix}_${cleanup_affix} \
  exp/tri3_reseg${affix}_$cleanup_affix{,/graph${affix}_${cleanup_affix}} 

for dset in dev test; do
  steps/decode_fmllr.sh --nj $decode_nj --cmd "$decode_cmd"  --num-threads 4 \
    exp/tri3_reseg${affix}_$cleanup_affix/graph${affix}_${cleanup_affix} \
    data/${dset} \
    exp/tri3_reseg${affix}_$cleanup_affix/decode${affix}_${cleanup_affix}_${dset}
  steps/lmrescore_const_arpa.sh  --cmd "$decode_cmd" data/lang${affix}_${cleanup_affix} \
    data/lang${affix}_${cleanup_affix}_rescore \
     data/${dset} exp/tri3_reseg${affix}_$cleanup_affix/decode${affix}_${cleanup_affix}_${dset} \
     exp/tri3_reseg${affix}_$cleanup_affix/decode${affix}_${cleanup_affix}_${dset}_rescore
done
) &

steps/get_prons.sh --cmd "$train_cmd" \
  data/train_reseg${affix} \
  data/lang_nosp exp/tri3_reseg${affix}
utils/dict_dir_add_pronprobs.sh --max-normalize true \
  data/local/dict_nosp \
  exp/tri3_reseg${affix}/{pron,sil,pron_bigram}_counts_nowb.txt \
  data/local/dict${affix}

utils/prepare_lang.sh data/local/dict${affix} \
  "<unk>" data/local/lang data/lang${affix}
cp -rT data/lang${affix} data/lang${affix}_rescore
cp data/lang_nosp/G.fst data/lang${affix}/
cp data/lang_nosp_rescore/G.carpa data/lang${affix}_rescore/

(
utils/mkgraph.sh data/lang${affix} \
  exp/tri3_reseg${affix}{,/graph${affix}} 

for dset in dev test; do
  steps/decode_fmllr.sh --nj $decode_nj --cmd "$decode_cmd"  --num-threads 4 \
    exp/tri3_reseg${affix}/graph${affix} \
    data/${dset} \
    exp/tri3_reseg${affix}/decode${affix}_${dset}
  steps/lmrescore_const_arpa.sh  --cmd "$decode_cmd" data/lang${affix} \
    data/lang${affix}_rescore \
     data/${dset} exp/tri3_reseg${affix}/decode${affix}_${dset} \
     exp/tri3_reseg${affix}/decode${affix}_${dset}_rescore
done
) &

wait
exit 0