#!/usr/bin/env bash

set -euo pipefail

# This script is called from local/nnet3/run_tdnn.sh and
# local/chain/run_tdnn.sh (and may eventually be called by more
# scripts).  It contains the common feature preparation and
# iVector-related parts of the script.  See those scripts for examples
# of usage.

stage=0
nj=56
num_threads_ubm=2

train_set=train
test_sets="native nonnative devtest test"
gmm=tri3b

nnet3_affix=

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

gmm_dir=exp/${gmm}
ali_dir=exp/${gmm}_ali_${train_set}_sp

for f in data/${train_set}/feats.scp ${gmm_dir}/final.mdl; do
  if [ ! -f $f ]; then
    echo "$0: expected file $f to exist"
    exit 1
  fi
done

if [ $stage -le 1 ]; then
    # perturb data to get alignments
    # nnet will be trained by high resolution data
    # _sp stands for speed-perturbed
    echo "$0: preparing directory for low-resolution speed-perturbed data (for alignment)"
    utils/data/perturb_data_dir_speed_3way.sh \
	data/${train_set} \
	data/${train_set}_sp

    echo "$0: making mfcc features for low-resolution speed-perturbed data"
    steps/make_mfcc.sh --cmd "$train_cmd" --nj 10 data/${train_set}_sp || exit 1;
    steps/compute_cmvn_stats.sh data/${train_set}_sp || exit 1;
    utils/fix_data_dir.sh data/${train_set}_sp
fi

if [ $stage -le 2 ]; then
    echo "$0: aligning with the perturbed low-resolution data"
    steps/align_fmllr.sh \
	--nj 20 --cmd "$train_cmd" data/${train_set}_sp data/lang $gmm_dir \
	$ali_dir || exit 1
fi

if [ $stage -le 3 ]; then
    # Create high-resolution MFCC features (with 40 cepstra instead of 13).

    echo "$0: creating high-resolution MFCC features"
    mfccdir=data/${train_set}_sp_hires/data
    for datadir in ${train_set}_sp ${test_sets}; do
	utils/copy_data_dir.sh \
	    data/$datadir \
	    data/${datadir}_hires
    done

    # do volume-perturbation on the training data prior to extracting hires
    # features; this helps make trained nnets more invariant to test data volume.
    utils/data/perturb_data_dir_volume.sh \
	data/${train_set}_sp_hires || exit 1;

    for datadir in ${train_set}_sp ${test_sets}; do
	steps/make_mfcc.sh \
	    --nj 10 \
	    --mfcc-config conf/mfcc_hires.conf \
	    --cmd "$train_cmd" \
	    data/${datadir}_hires || exit 1;
	steps/compute_cmvn_stats.sh \
	    data/${datadir}_hires || exit 1;
	utils/fix_data_dir.sh \
	    data/${datadir}_hires || exit 1;
    done
fi

if [ $stage -le 4 ]; then
    echo "$0: computing a subset of data to train the diagonal UBM."
    # We'll use about a quarter of the data.
    mkdir -p exp/nnet3${nnet3_affix}/diag_ubm
    temp_data_root=exp/nnet3${nnet3_affix}/diag_ubm

    num_utts_total=$(wc -l <data/${train_set}_sp_hires/utt2spk)
    num_utts=$[$num_utts_total/4]
    utils/data/subset_data_dir.sh \
	data/${train_set}_sp_hires \
	$num_utts \
	${temp_data_root}/${train_set}_sp_hires_subset

    echo "$0: computing a PCA transform from the hires data."
    steps/online/nnet2/get_pca_transform.sh \
	--cmd "$train_cmd" \
	--splice-opts "--left-context=3 --right-context=3" \
	--max-utts 10000 \
	--subsample 2 \
	${temp_data_root}/${train_set}_sp_hires_subset \
	exp/nnet3${nnet3_affix}/pca_transform

    echo "$0: training the diagonal UBM."
    # Use 512 Gaussians in the UBM.
    steps/online/nnet2/train_diag_ubm.sh \
	--cmd "$train_cmd" \
	--nj 20 \
	--num-frames 700000 \
	--num-threads 8 \
	${temp_data_root}/${train_set}_sp_hires_subset \
	512 \
	exp/nnet3${nnet3_affix}/pca_transform \
	exp/nnet3${nnet3_affix}/diag_ubm
fi

if [ $stage -le 5 ]; then
    # Train the iVector extractor.
    # Use all the speed-perturbed data .
    # iVector extractors can be sensitive to the amount of data.
    # The script defaults to an iVector dimension of 100.
    echo "$0: training the iVector extractor"
    steps/online/nnet2/train_ivector_extractor.sh \
	--cmd "$train_cmd" \
	--nj 10 \
	data/${train_set}_sp_hires \
	exp/nnet3${nnet3_affix}/diag_ubm \
	exp/nnet3${nnet3_affix}/extractor || exit 1;
fi

if [ $stage -le 6 ]; then
    # combine   and train system on short segments.
    # extract iVectors on speed-perturbed training data
    # With --utts-per-spk-max 2, script pairs  utterances into twos.
    # Treats each  pair as one speaker.
    # Gives more diversity in iVectors.
    # Extracted online.

    # note: extract  ivectors from max2 data
    # Why is max2 not encoded in ivectordir name?
    # valid for non-max2 data
    #  utterance list is the same.

    ivectordir=exp/nnet3${nnet3_affix}/ivectors_${train_set}_sp_hires

    # having a larger number of speakers is helpful for generalization, and to
    # handle per-utterance decoding well (iVector starts at zero).
    temp_data_root=${ivectordir}
    utils/data/modify_speaker_info.sh \
	--utts-per-spk-max 2 \
	data/${train_set}_sp_hires \
	${temp_data_root}/${train_set}_sp_hires_max2

    steps/online/nnet2/extract_ivectors_online.sh \
	--cmd "$train_cmd" \
	--nj 20 \
	${temp_data_root}/${train_set}_sp_hires_max2 \
	exp/nnet3${nnet3_affix}/extractor \
	$ivectordir

    # Also extract iVectors for test data.
    # No need for speed perturbation (sp).
    for data in $test_sets; do
	steps/online/nnet2/extract_ivectors_online.sh \
	    --cmd "$train_cmd" \
	    --nj 8 \
	    data/${data}_hires \
	    exp/nnet3${nnet3_affix}/extractor \
	    exp/nnet3${nnet3_affix}/ivectors_${data}_hires
    done
fi

exit 0
