#!/bin/bash
  
  # Copyright  2017  Nagendra Kumar Goel
  #            2017  Vimal Manohar
  # Apache 2.0
  
  # We assume the run.sh has been executed (because we are using model
  # directories like exp/tri4a)
  
  # This script demonstrates nnet3-based speech activity detection for
  # segmentation.
  # This script:
  # 1) Prepares targets (per-frame labels) for a subset of training data 
  #    using GMM models
  # 2) Augments the training data with reverberation and additive noise
  # 3) Trains TDNN+Stats or TDNN+LSTM neural network using the targets 
  #    and augmented data
  # 4) Demonstrates using the SAD system to get segments of dev data and decode
  
  lang=data/lang   # Must match the one used to train the models
  lang_test=data/lang_test  # Lang directory for decoding.
  
  data_dir=data/train_100k
  # Model directory used to align the $data_dir to get target labels for training
  # SAD. This should typically be a speaker-adapted system.
  sat_model_dir=exp/tri4a
  # Model direcotry used to decode the whole-recording version of the $data_dir to
  # get target labels for training SAD. This should typically be a
  # speaker-independent system like LDA+MLLT system.
  model_dir=exp/tri3a
  graph_dir=exp/tri3a/graph   # Graph for decoding whole-recording version of $data_dir.
                              # If not provided, a new one will be created using $lang_test
  
  # List of weights on labels obtained from alignment;
  # labels obtained from decoding; and default labels in out-of-segment regions
  merge_weights=1.0,0.1,0.5
  
  prepare_targets_stage=-10
  nstage=-10
  train_stage=-10
  test_stage=-10
  num_data_reps=3
  affix=_1a   # For segmentation
  test_affix=1a
  stage=-1
  nj=80
  reco_nj=40
  
  # test options
  test_nj=30
  
  . ./cmd.sh
  if [ -f ./path.sh ]; then . ./path.sh; fi
  
  set -e -u -o pipefail
  . utils/parse_options.sh 
  
  if [ $# -ne 0 ]; then
    exit 1
  fi
  
  dir=exp/segmentation${affix}
  mkdir -p $dir
  
  # See $lang/phones.txt and decide which should be garbage
  garbage_phones="laughter oov"
  silence_phones="sil noise"
  
  for p in $garbage_phones; do 
    for a in "" "_B" "_E" "_I" "_S"; do
      echo "$p$a"
    done
  done > $dir/garbage_phones.txt
  
  for p in $silence_phones; do 
    for a in "" "_B" "_E" "_I" "_S"; do
      echo "$p$a"
    done
  done > $dir/silence_phones.txt
  
  if ! cat $dir/garbage_phones.txt $dir/silence_phones.txt | \
    steps/segmentation/internal/verify_phones_list.py $lang/phones.txt; then
    echo "$0: Invalid $dir/{silence,garbage}_phones.txt"
    exit 1
  fi
  
  whole_data_dir=${data_dir}_whole
  whole_data_id=$(basename $whole_data_dir)
  
  rvb_data_dir=${whole_data_dir}_rvb_hires
  
  if [ $stage -le 0 ]; then
    utils/data/convert_data_dir_to_whole.sh $data_dir $whole_data_dir
  fi
  
  ###############################################################################
  # Extract features for the whole data directory
  ###############################################################################
  if [ $stage -le 1 ]; then
    steps/make_mfcc.sh --nj $reco_nj --cmd "$train_cmd"  --write-utt2num-frames true \
      $whole_data_dir exp/make_mfcc/${whole_data_id}
    steps/compute_cmvn_stats.sh $whole_data_dir exp/make_mfcc/${whole_data_id}
    utils/fix_data_dir.sh $whole_data_dir
  fi
  
  ###############################################################################
  # Prepare SAD targets for recordings
  ###############################################################################
  targets_dir=$dir/${whole_data_id}_combined_targets_sub3
  if [ $stage -le 3 ]; then
    steps/segmentation/prepare_targets_gmm.sh --stage $prepare_targets_stage \
      --train-cmd "$train_cmd" --decode-cmd "$decode_cmd" \
      --nj $nj --reco-nj $reco_nj --lang-test $lang_test \
      --garbage-phones-list $dir/garbage_phones.txt \
      --silence-phones-list $dir/silence_phones.txt \
      --merge-weights "$merge_weights" \
      --graph-dir "$graph_dir" \
      $lang $data_dir $whole_data_dir $sat_model_dir $model_dir $dir
  fi
  
  rvb_targets_dir=${targets_dir}_rvb
  if [ $stage -le 4 ]; then
    # Download the package that includes the real RIRs, simulated RIRs, isotropic noises and point-source noises
    if [ ! -f rirs_noises.zip ]; then
      wget --no-check-certificate http://www.openslr.org/resources/28/rirs_noises.zip
      unzip rirs_noises.zip
    fi
  
    rvb_opts=()
    # This is the config for the system using simulated RIRs and point-source noises
    rvb_opts+=(--rir-set-parameters "0.5, RIRS_NOISES/simulated_rirs/smallroom/rir_list")
    rvb_opts+=(--rir-set-parameters "0.5, RIRS_NOISES/simulated_rirs/mediumroom/rir_list")
    rvb_opts+=(--noise-set-parameters RIRS_NOISES/pointsource_noises/noise_list)
  
    foreground_snrs="20:10:15:5:0"
    background_snrs="20:10:15:5:0"
    # corrupt the data to generate multi-condition data
    # for data_dir in train dev test; do
    python steps/data/reverberate_data_dir.py \
      "${rvb_opts[@]}" \
      --prefix "rev" \
      --foreground-snrs $foreground_snrs \
      --background-snrs $background_snrs \
      --speech-rvb-probability 0.5 \
      --pointsource-noise-addition-probability 0.5 \
      --isotropic-noise-addition-probability 0.7 \
      --num-replications $num_data_reps \
      --max-noises-per-minute 4 \
      --source-sampling-rate 8000 \
      $whole_data_dir $rvb_data_dir
  fi
  
  if [ $stage -le 5 ]; then
    steps/make_mfcc.sh --mfcc-config conf/mfcc_hires.conf --nj $reco_nj \
      ${rvb_data_dir}
    steps/compute_cmvn_stats.sh ${rvb_data_dir}
    utils/fix_data_dir.sh $rvb_data_dir
  fi
  
  if [ $stage -le 6 ]; then
    rvb_targets_dirs=()
    for i in `seq 1 $num_data_reps`; do
      steps/segmentation/copy_targets_dir.sh --utt-prefix "rev${i}_" \
        $targets_dir ${targets_dir}_temp_$i || exit 1
      rvb_targets_dirs+=(${targets_dir}_temp_$i)
    done
  
    steps/segmentation/combine_targets_dirs.sh \
      $rvb_data_dir ${rvb_targets_dir} \
      ${rvb_targets_dirs[@]} || exit 1;
  
    rm -r ${rvb_targets_dirs[@]}
  fi
  
  
  sad_nnet_dir=$dir/tdnn_stats_asr_sad_1a
  
  if [ $stage -le 7 ]; then
    # Train a STATS-pooling network for SAD
    local/segmentation/tuning/train_stats_asr_sad_1a.sh \
      --stage $nstage --train-stage $train_stage \
      --targets-dir ${rvb_targets_dir} \
      --data-dir ${rvb_data_dir} --affix "1a" || exit 1
  
    # # Train a TDNN+LSTM network for SAD
    # local/segmentation/tuning/train_lstm_asr_sad_1a.sh \
    #   --stage $nstage --train-stage $train_stage \
    #   --targets-dir ${rvb_targets_dir} \
    #   --data-dir ${rvb_data_dir} --affix "1a" || exit 1
  fi
  
  if [ ! -f data/dev_aspire/wav.scp ]; then
    echo "$0: Not evaluating on data/dev_aspire"
    exit 0
  fi
  
  if [ $stage -le 8 ]; then
  steps/segmentation/convert_utt2spk_and_segments_to_rttm.py \
    --reco2file-and-channel=data/dev_aspire/reco2file_and_channel \
    data/dev_aspire/{utt2spk,segments,ref.rttm}
  fi
  
  chain_dir=exp/chain/tdnn_lstm_1a
  
  # The context options in "sad_opts" must match the options used to train the 
  # SAD network in "sad_nnet_dir"
  sad_opts="--extra-left-context 79 --extra-right-context 21 --frames-per-chunk 150 --extra-left-context-initial 0 --extra-right-context-final 0 --acwt 0.3"
  
  # For LSTM SAD network, the options might be something like
  # sad_opts="--extra-left-context 70 --extra-right-context 0 --frames-per-chunk 150 --extra-left-context-initial 0 --extra-right-context-final 0 --acwt 0.3"
  
  if [ $stage -le 9 ]; then
    # Use left and right context options that were used when training
    # the chain nnet
    # Increase sil-scale to predict more silence
    local/nnet3/prep_test_aspire_segmentation.sh --stage $test_stage \
      --decode-num-jobs $test_nj --affix "${test_affix}" \
      --sad-opts "$sad_opts" \
      --sad-graph-opts "--min-silence-duration=0.03 --min-speech-duration=0.3 --max-speech-duration=10.0" --sad-priors-opts "--sil-scale=0.1" \
      --acwt 1.0 --post-decode-acwt 10.0 \
      --extra-left-context 50 \
      --extra-right-context 0 \
      --extra-left-context-initial 0 --extra-right-context-final 0 \
     --sub-speaker-frames 6000 --max-count 75 \
     --decode-opts "--min-active 1000" \
     dev_aspire $sad_nnet_dir $sad_nnet_dir data/lang $chain_dir/graph_pp $chain_dir
  fi