#!/bin/bash

# Copyright 2014  Vimal Manohar
# This is our online neural net build for Gale system

. ./cmd.sh

stage=-1
train_stage=-10
use_gpu=true
mfccdir=mfcc
train_nj=120
decode_nj=30

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

if $use_gpu; then
  if ! cuda-compiled; then
    cat <<EOF && exit 1
This script is intended to be used with GPUs but you have not compiled Kaldi with CUDA
If you want to use GPUs (and have them), go to src/, and configure and make on a machine
where "nvcc" is installed.  Otherwise, call this script with --use-gpu false
EOF
  fi
  parallel_opts="--gpu 1"
  num_threads=1
  minibatch_size=512
  # the _a is in case I want to change the parameters.
  dir=exp/nnet2_online/nnet_a_gpu
else
  # Use 4 nnet jobs just like run_4d_gpu.sh so the results should be
  # almost the same, but this may be a little bit slow.
  num_threads=16
  minibatch_size=128
  parallel_opts="--num-threads $num_threads"
  dir=exp/nnet2_online/nnet_a
fi

if [ $stage -le 0 ]; then
  # this shows how you can split across multiple file-systems.  we'll split the
  # MFCC dir across multiple locations.  You might want to be careful here, if you
  # have multiple copies of Kaldi checked out and run the same recipe, not to let
  # them overwrite each other.
  if [[ $(hostname -f) == *.clsp.jhu.edu ]] && [ ! -d $mfccdir/storage ]; then
    date=$(date +'%m_%d_%H_%M')
    utils/create_split_dir.pl /export/b0{1,2,3,4}/$USER/kaldi-data/egs/gale-$date/s5/$mfccdir/storage $mfccdir/storage || exit 1
  fi
  utils/copy_data_dir.sh data/train data/train_hires || exit 1
  steps/make_mfcc_pitch_online.sh --nj $train_nj --mfcc-config conf/mfcc_hires.conf \
      --cmd "$train_cmd" data/train_hires exp/make_hires/train $mfccdir || exit 1;
  steps/compute_cmvn_stats.sh data/train_hires exp/make_hires/train $mfccdir || exit 1;
fi

if [ $stage -le 1 ]; then
  # we'll use the features with just MFCC, no pitch, to train the iVector
  # extractor on.  Check that we're using 40-dim features so the command line is correct.
  ! grep 'num-ceps=40' conf/mfcc_hires.conf >/dev/null && \
     echo "Change the script if you change conf/mfcc_hires.conf" && exit 1;
  steps/select_feats.sh  --nj 5 --cmd "$train_cmd" 0-39 data/train_hires \
      data/train_hires_mfcconly exp/nnet2_online/select_hires_train $mfccdir || exit 1

  steps/compute_cmvn_stats.sh data/train_hires_mfcconly exp/nnet2_online/select_hires_train $mfccdir || exit 1

  # Make a subset of about 1/3 of the data.
  utils/subset_data_dir.sh data/train_hires_mfcconly 100000 \
     data/train_hires_mfcconly_100k || exit 1

  # make a corresponding subset of normal-dimensional-MFCC training data.
  utils/subset_data_dir.sh --utt-list <(awk '{print $1}' data/train_hires_mfcconly_100k/utt2spk) \
    data/train data/train_100k || exit 1
fi

if [ $stage -le 2 ]; then
  # We need to build a small system just because we need the LDA+MLLT transform
  # to train the diag-UBM on top of.  First align the data of the 100k subset using
  # the tri3b system and normal MFCC features, so we have alignments to build our
  # system on hires MFCCs on top of.

  steps/align_fmllr.sh --nj $train_nj --cmd "$train_cmd" \
    data/train_100k data/lang exp/tri3b exp/tri3b_ali_100k || exit 1;

  # Build a small LDA+MLLT system on top of the hires MFCC features, just
  # because we need the transform.  We use --num-iters 13 because after we get
  # the transform (12th iter is the last), any further training is pointless.
  steps/train_lda_mllt.sh --cmd "$train_cmd" --num-iters 13 --realign-iters "" \
    --splice-opts "--left-context=3 --right-context=3" \
    5000 10000 data/train_hires_mfcconly_100k data/lang exp/tri3b_ali_100k exp/nnet2_online/tri4a || exit 1
fi

if [ $stage -le 3 ]; then
  # Train a diagonal UBM.  The input directory exp/nnet2_online/tri3a is only
  # needed for the splice-opts and the LDA+MLLT transform.
  steps/online/nnet2/train_diag_ubm.sh --cmd "$train_cmd" --nj $train_nj --num-frames 400000 \
    data/train_hires_mfcconly_100k 512 exp/nnet2_online/tri4a exp/nnet2_online/diag_ubm || exit 1
fi

if [ $stage -le 4 ]; then
  # train an iVector extractor on all the mfcconly data.  Note: although we use
  # only 10 job, each job uses 16 processes in total.
  steps/online/nnet2/train_ivector_extractor.sh --cmd "$train_cmd" --nj 10 \
    data/train_hires_mfcconly exp/nnet2_online/diag_ubm exp/nnet2_online/extractor || exit 1;
fi

if [ $stage -le 5 ]; then
  # extract iVectors for the training data.
  ivectordir=exp/nnet2_online/ivectors_train
  if [[ $(hostname -f) == *.clsp.jhu.edu ]] && [ ! -d $ivectordir/storage ]; then # this shows how you can split across multiple file-systems.
    utils/create_split_dir.pl /export/b0{1,2,3,4}/$USER/kaldi-data/egs/gale/s5/$ivectordir/storage $ivectordir/storage || exit 1
  fi

  # having a larger number of speakers is helpful for generalization, and to
  # handle per-utterance decoding well (iVector starts at zero).
  steps/online/nnet2/copy_data_dir.sh --utts-per-spk-max 2 data/train_hires_mfcconly data/train_hires_mfcconly_max2 || exit 1

  steps/online/nnet2/extract_ivectors_online.sh --cmd "$train_cmd" --nj $train_nj \
    data/train_hires_mfcconly_max2 exp/nnet2_online/extractor $ivectordir || exit 1;
fi

if [ $stage -le 6 ]; then
  # this shows how you can split across multiple file-systems.
  if [[ $(hostname -f) == *.clsp.jhu.edu ]] && [ ! -d $dir/egs/storage ]; then
    utils/create_split_dir.pl /export/b0{1,2,3,4}/$USER/kaldi-online/egs/bolt/s5/$dir/egs $dir/egs/storage || exit 1
  fi

  # Because we have a lot of data here and we don't want the training to take
  # too long, we reduce the number of epochs from the defaults (15) to (8).
  # The option "--io-opts '--max-jobs-run 12'" is to have more than the default number
  # (5) of jobs dumping the egs to disk; this is OK since we're splitting our
  # data across four filesystems for speed.

  steps/nnet2/train_pnorm_simple.sh --stage $train_stage \
    --num-epochs 8 \
    --samples-per-iter 400000 \
    --splice-width 7 --feat-type raw \
    --online-ivector-dir exp/nnet2_online/ivectors_train \
    --cmvn-opts "--norm-means=false --norm-vars=false" \
    --num-threads "$num_threads" \
    --minibatch-size "$minibatch_size" \
    --parallel-opts "$parallel_opts" \
    --io-opts "--max-jobs-run 12" \
    --num-jobs-nnet 6 \
    --num-hidden-layers 4 \
    --mix-up 12000 \
    --initial-learning-rate 0.06 --final-learning-rate 0.006 \
    --cmd "$decode_cmd" \
    --pnorm-input-dim 3000 \
    --pnorm-output-dim 300 \
     data/train_hires data/lang exp/tri3b $dir  || exit 1;
fi

if [ $stage -le 7 ]; then
  steps/online/nnet2/prepare_online_decoding.sh --mfcc-config conf/mfcc_hires.conf \
     --add-pitch true data/lang exp/nnet2_online/extractor "$dir" ${dir}_online || exit 1;
fi

if [ $stage -le 8 ]; then
  # do the actual online decoding with iVectors, carrying info forward from
  # previous utterances of the same speaker.
   steps/online/nnet2/decode.sh --config conf/decode.config --cmd "$decode_cmd" --nj $decode_nj \
      exp/tri3b/graph data/test ${dir}_online/decode_test || exit 1;
fi

if [ $stage -le 9 ]; then
  # this version of the decoding treats each utterance separately
  # without carrying forward speaker information.
   steps/online/nnet2/decode.sh --config conf/decode.config --cmd "$decode_cmd" --nj $decode_nj \
     --per-utt true \
      exp/tri3b/graph data/test ${dir}_online/decode_test_utt || exit 1;
fi

if [ $stage -le 10 ]; then
  # this version of the decoding treats each utterance separately
  # without carrying forward speaker information, but looks to the end
  # of the utterance while computing the iVector.
   steps/online/nnet2/decode.sh --config conf/decode.config --cmd "$decode_cmd" --nj $decode_nj \
     --per-utt true --online false \
      exp/tri3b/graph data/test ${dir}_online/decode_test_utt_offline || exit 1;
fi

exit 0;