#!/bin/bash

# Kaldi ASR baseline for the CHiME-4 Challenge (2ch track: 2 channel track)
#
# Copyright 2016 University of Sheffield (Jon Barker, Ricard Marxer)
#                Inria (Emmanuel Vincent)
#                Mitsubishi Electric Research Labs (Shinji Watanabe)
#           2017 JHU CLSP (Szu-Jui Chen)
#           2017 JHU CLSP (Aswin Shanmugam Subramanian)
#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)

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

#####Baseline settings#####
# Usage: 
# Execute './run.sh' to get the models.
# We provide three kinds of beamform methods. Add option --enhancement blstm_gev, or --enhancement beamformit_2mics
# to use them. i.g. './run.sh --enhancement blstm_gev'
#
# We stopped to support the old CHiME-3/4 baseline. If you want to reproduce the old results
# Please use the old version of Kaldi, e.g., git checkout 9e8ff73648917836d0870c8f6fdd2ff4bdde384f

# Config:
stage=0 # resume training with --stage N
enhancement=blstm_gev #### or your method 
# if the following options are true, they wouldn't train a model again and will only do decoding
gmm_decode_only=false
tdnn_decode_only=false
# make it true when you want to add enhanced data into training set. But please note that when changing enhancement method,
# you may need to retrain from run_gmm.sh and avoid using decode-only options above
add_enhanced_data=true

. utils/parse_options.sh || exit 1;

# Set bash to 'debug' mode, it will exit on :
# -e 'error', -u 'undefined variable', -o ... 'error in pipeline', -x 'print commands',
set -e
set -u
set -o pipefail

#####check data and model paths################
# Set a main root directory of the CHiME4 data
# If you use scripts distributed in the CHiME4 package,
chime4_data=`pwd`/../..
# Otherwise, please specify it, e.g.,
# chime4_data=/db/laputa1/data/processed/public/CHiME4
# chime3_data=/data2/archive/speech-db/original/public/CHiME3

case $(hostname -f) in
  *.clsp.jhu.edu) 
      chime4_data=/export/corpora4/CHiME4/CHiME3 # JHU,
      chime3_data=/export/corpora5/CHiME3 
      ;;
esac 

if [ ! -d $chime4_data ]; then
  echo "$chime4_data does not exist. Please specify chime4 data root correctly" && exit 1;
fi
if [ ! -d $chime3_data ]; then
  echo "$chime3_data does not exist. Please specify chime4 data root correctly" && exit 1;
fi

#####main program start################
# You can execute run_init.sh only "once"
# This creates 3-gram LM, FSTs, and basic task files
if [ $stage -le 0 ]; then
  local/run_init.sh $chime4_data
fi

# Using Beamformit or mask-based beamformer
# note that beamformed WAV files are generated in the following directory
enhancement_data=`pwd`/enhan/$enhancement
if [ $stage -le 1 ]; then
   case $enhancement in
    beamformit_2mics)
        local/run_beamform_2ch_track.sh --cmd "$train_cmd" --nj 20 $chime4_data/data/audio/16kHz/isolated_2ch_track $enhancement_data
        ;;
    blstm_gev)
        local/run_blstm_gev.sh --cmd "$train_cmd" --nj 20 --track 2 $chime4_data $chime3_data $enhancement_data 0
        ;;
    *)
        echo "Usage: --enhancement blstm_gev, or --enhancement beamformit_2mics" 
        exit 1;
   esac
fi

# Compute PESQ, STOI, eSTOI, and SDR scores
if [ $stage -le 2 ]; then
  if [ ! -f local/bss_eval_sources.m ] || [ ! -f local/stoi.m ] || [ ! -f local/estoi.m ] || [ ! -f local/PESQ ]; then
    # download and install speech enhancement evaluation tools
    local/download_se_eval_tool.sh
  fi
  chime4_rir_data=local/nn-gev/data/audio/16kHz/isolated_ext
  if [ ! -d $chime4_rir_data ]; then
    echo "$chime4_rir_dir does not exist. Please run 'blstm_gev' enhancement method first;" && exit 1;
  fi
  local/compute_pesq.sh $enhancement $enhancement_data $chime4_rir_data $PWD
  local/compute_stoi_estoi_sdr.sh $enhancement $enhancement_data $chime4_rir_data
  local/compute_pesq.sh NOISY_1ch $chime4_data/data/audio/16kHz/isolated_1ch_track/ $chime4_rir_data $PWD
  local/compute_stoi_estoi_sdr.sh NOISY_1ch $chime4_data/data/audio/16kHz/isolated_1ch_track/ $chime4_rir_data
  local/write_se_results.sh $enhancement
  local/write_se_results.sh NOISY_1ch
fi

# GMM based ASR experiment
# Please set a directory of your speech enhancement method.
# The directory structure and audio files must follow the attached baseline enhancement directory
if [ $stage -le 3 ]; then
  local/run_gmm.sh --add-enhanced-data $add_enhanced_data \
    --decode-only $gmm_decode_only $enhancement $enhancement_data $chime4_data
fi

# TDNN based ASR experiment
# Since it takes time to evaluate TDNN, we make the GMM and TDNN scripts separately.
# You may execute it after you would have promising results using GMM-based ASR experiments
if [ $stage -le 4 ]; then
  local/chain/run_tdnn.sh --decode-only $tdnn_decode_only $enhancement
fi

# LM-rescoring experiment with 5-gram and RNN LMs
# It takes a few days to train a RNNLM.
if [ $stage -le 5 ]; then
  local/run_lmrescore_tdnn.sh $chime4_data $enhancement
fi

# LM-rescoring experiment with LSTM LMs
if [ $stage -le 6 ]; then
  local/rnnlm/run_lstm.sh $enhancement
fi

echo "Done."