#!/bin/bash
  
  # Copyright 2016 Pegah Ghahremani
  
  # This script is used for decoding multilingual model and it is called in
  # local/nnet3/run_tdnn_multilingual.sh script.
  # This script needs decoding data dir, which is prepared using
  # eg/babel/s5d scripts (i.e. run-4-anydecode.sh).
  # If --use-pitch is true, pitch feature is added to high-resolution MFCC features.
  # If --use-bnf option is true, the --bnf-nnet-dir option, nnet3 model for
  # bottleneck feature extraction, should be provided.
  
  set -e
  set -o pipefail
  
  
  dir=dev10h.pem
  kind=
  use_pitch=true
  use_pitch_ivector=false # If true, pitch feature is used in ivector extraction.
  use_ivector=false
  decode_stage=-1
  nnet3_affix=
  feat_suffix=
  ivector_suffix=
  iter=final
  nj=30
  
  # params for extracting bn features
  use_bnf=false # If true, bottleneck feature is extracted and appended to input
                # for nnet3 model.
  bnf_nnet_dir=exp/nnet3/multi_bnf_sp # dir for bottlneck nnet3 model
                                      # (used for bottleneck feature extraction)
  use_ivector_bnf=false # If true, ivector used in extracting bottleneck features.
  
  . conf/common_vars.sh || exit 1;
  
  . utils/parse_options.sh
  
  if [ $# -ne 2 ]; then
    echo "Usage: $(basename $0) --dir <dir-type> <lang> <multilingual-nnet3-dir>"
    echo " e.g.: $(basename $0) --dir dev2h.pem ASM exp/nnet3/tdnn_multi_sp"
    exit 1
  fi
  
  lang=$1
  nnet3_dir=$2
  
  langconf=conf/$lang/lang.conf
  
  if [ ! -f $langconf ]; then
    echo "$0: Language configuration $langconf does not exist! Use the "
    echo "configurations in ../../babel/s5d/conf/lang/$lang-* as a startup." && exit 1
  fi
  . $langconf || exit 1;
  [ -f local.conf ] && . local.conf;
  
  mfcc=mfcc/$lang
  data=data/$lang
  vector_suffix=_gb
  
  dataset_dir=$data/$dir
  dataset_id=$dir
  dataset_type=${dir%%.*}
  
  #By default, we want the script to accept how the dataset should be handled,
  #i.e. of  what kind is the dataset
  if [ -z ${kind} ] ; then
    if [ "$dataset_type" == "dev2h" ] || [ "$dataset_type" == "dev10h" ]; then
      dataset_kind=supervised
    else
      dataset_kind=unsupervised
    fi
  else
    dataset_kind=$kind
  fi
  
  dataset=$(basename $dataset_dir)
  mfccdir=mfcc_hires/$lang
  mfcc_affix=""
  hires_config="--mfcc-config conf/mfcc_hires.conf"
  nnet3_data_dir=${dataset_dir}_hires
  feat_suffix=_hires
  ivec_feat_suffix=_hires
  log_dir=exp/$lang/make_hires/$dataset
  
  if $use_pitch_ivector; then
    ivec_feat_suffix=_hires_pitch
  fi
  
  if $use_pitch; then
    mfcc_affix="_pitch_online"
    hires_config="$hires_config --online-pitch-config conf/pitch.conf"
    mfccdir=mfcc_hires_pitch/lang
    nnet3_data_dir=${dataset_dir}_hires_pitch
    feat_suffix="_hires_pitch"
    log_dir=exp/$lang/make_hires_pitch/$dataset
  fi
  
  
  ####################################################################
  ##
  ##  Feature extraction for decoding
  ##
  ####################################################################
  echo ---------------------------------------------------------------------
  echo "Preparing ${dataset_kind} data files in ${dataset_dir} on" `date`
  echo ---------------------------------------------------------------------
  if [ ! -f  $dataset_dir/.done ] ; then
    if [ ! -f ${nnet3_data_dir}/.mfcc.done ]; then
      echo ---------------------------------------------------------------------
      echo "Preparing ${dataset_kind} MFCC features in  ${nnet3_data_dir} and corresponding "
      echo "iVectors in exp/$lang/nnet3${nnet3_affix}/ivectors_${dataset}${feat_suffix}${ivector_suffix} on" `date`
      echo ---------------------------------------------------------------------
      if [ ! -d ${nnet3_data_dir} ]; then
        utils/copy_data_dir.sh $data/$dataset ${nnet3_data_dir}
      fi
  
      steps/make_mfcc${mfcc_affix}.sh --nj $nj $hires_config \
          --cmd "$train_cmd" ${nnet3_data_dir} $log_dir $mfccdir;
      steps/compute_cmvn_stats.sh ${nnet3_data_dir} $log_dir $mfccdir;
      utils/fix_data_dir.sh ${nnet3_data_dir};
      touch ${nnet3_data_dir}/.mfcc.done
    fi
    touch $dataset_dir/.done
  fi
  
  ivector_dir=exp/$lang/nnet3${nnet3_affix}/ivectors_${dataset}${ivec_feat_suffix}${ivector_suffix}
  if $use_ivector && [ ! -f $ivector_dir/.ivector.done ];then
    extractor=exp/multi/nnet3${nnet3_affix}/extractor
    ivec_feat_suffix=$feat_suffix
    if $use_pitch && ! $use_pitch_ivector; then
      ivec_feat_suffix=_hires
      featdir=${dataset_dir}${feat_suffix}
      mfcc_only_dim=`feat-to-dim scp:$featdir/feats.scp - | awk '{print $1-3}'`
      steps/select_feats.sh --cmd "$train_cmd" --nj $nj 0-$[$mfcc_only_dim-1] \
        $featdir ${dataset_dir}${ivec_feat_suffix} || exit 1;
      steps/compute_cmvn_stats.sh ${dataset_dir}${ivec_feat_suffix} || exit 1;
    fi
  
    steps/online/nnet2/extract_ivectors_online.sh --cmd "$train_cmd" --nj $nj \
      ${dataset_dir}${ivec_feat_suffix} $extractor $ivector_dir || exit 1;
    touch $ivector_dir/.ivector.done
  fi
  
  if $use_bnf; then
    multi_ivector_dir=exp/$lang/nnet3${nnet3_affix}/ivectors_${dataset}${ivec_feat_suffix}${ivector_suffix}
  
    ivector_for_bnf_opt=
    if $use_ivector_bnf;then ivector_for_bnf_opt="--ivector-dir $multi_ivector_dir"; fi
  
    bnf_data_dir=${dataset_dir}_bnf/$lang
    if [ ! -f $bnf_data_dir/.done ]; then
      steps/nnet3/make_bottleneck_features.sh --use-gpu true --nj 100 --cmd "$train_cmd" \
        $ivector_for_bnf_opt tdnn_bn.renorm \
        ${dataset_dir}${feat_suffix} $bnf_data_dir \
        $bnf_nnet_dir bnf/$lang exp/$lang/make_${dataset}_bnf || exit 1;
      touch $bnf_data_dir/.done
    else
      echo "$0: Skip Bottleneck feature extraction; You can force to run this step deleting $bnf_data_dir/.done."
    fi
  
    appended_bnf=${dataset_dir}${feat_suffix}_bnf
    if [ ! -f $appended_bnf/.done ]; then
      steps/append_feats.sh  --nj 16 --cmd "$train_cmd" \
        $bnf_data_dir ${dataset_dir}${feat_suffix} \
        ${dataset_dir}${feat_suffix}_bnf exp/$lang/append${feat_suffix}_bnf \
        mfcc${feat_suffix}_bnf/$lang || exit 1;
  
      steps/compute_cmvn_stats.sh $appended_bnf exp/$lang/make_cmvn${feat_suffix}_bnf \
        mfcc${feat_suffix}_bnf/$lang || exit 1;
      touch $appended_bnf/.done
    fi
    feat_suffix=${feat_suffix}_bnf
  fi
  
  ####################################################################
  ##
  ## nnet3 model decoding
  ##
  ####################################################################
  if [ ! -f exp/$lang/tri5/graph/HCLG.fst ];then
    utils/mkgraph.sh \
      data/$lang/lang exp/$lang/tri5 exp/$lang/tri5/graph |tee exp/$lang/tri5/mkgraph.log
  fi
  
  if [ -f $nnet3_dir/$lang/final.mdl ]; then
    decode=$nnet3_dir/$lang/decode_${dataset_id}
    feat_suffix=_hires
    ivec_feat_suffix=_hires
  
    # suffix for using other features such as pitch
    if $use_pitch; then
      feat_suffix=${feat_suffix}_pitch
    fi
    if $use_pitch_ivector; then
      ivec_feat_suffix=_hires_pitch
    fi
    if $use_bnf; then
      feat_suffix=${feat_suffix}_bnf
    fi
    ivector_opts=
    if $use_ivector; then
      ivector_opts="--online-ivector-dir exp/$lang/nnet3${nnet3_affix}/ivectors_${dataset_id}${ivec_feat_suffix}${ivector_suffix}"
    fi
    if [ ! -f $decode/.done ]; then
      mkdir -p $decode
      score_opts="--skip-scoring false"
      [ ! -z $iter ] && iter_opt="--iter $iter"
      steps/nnet3/decode.sh --nj $nj --cmd "$decode_cmd" $iter_opt \
            --stage $decode_stage \
            --beam $dnn_beam --lattice-beam $dnn_lat_beam \
            $score_opts $ivector_opts \
            exp/$lang/tri5/graph ${dataset_dir}${feat_suffix} $decode | tee $decode/decode.log
  
      touch $decode/.done
    fi
  fi
  
  echo "Everything looking good...."
  exit 0