#!/bin/bash
#
# This script generates separate egs directory for each input
# language in multilingual setup, which contains both egs.*.ark and egs.*.scp.
#
# This script will generally be called from nnet3 multilingual training script.

echo "$0 $@"  # Print the command line for logging
set -e


# Begin configuration section
cmd=
stage=0
left_context=13
right_context=9
online_multi_ivector_dirs=     # list of iVector dir for all languages
                              # can be used if we are including speaker information as iVectors.
                              # e.g. "exp/lang1/train-ivector exp/lang2/train-ivector"
samples_per_iter=400000 # this is the target number of egs in each archive of egs
                        # (prior to merging egs).  We probably should have called
                        # it egs_per_iter. This is just a guideline; it will pick
                        # a number that divides the number of samples in the
                        # entire data.
# Configuration to allocate egs
minibatch_size=512
num_archives=100
num_jobs=10
cmvn_opts=
echo "$0 $@"  # Print the command line for logging

if [ -f path.sh ]; then . ./path.sh; fi
. parse_options.sh || exit 1;

if [ $# -lt 4 ]; then
  echo "Usage: $0 [opts] N <data-dir1> .. <data-dirN> <ali-dir1> .. <ali-dirN>"
  echo " <egs-out1> .. <egs-outN>"
  echo " e.g.: $0 2 data/lang1/train data/lang2/train exp/lang1/tri5_ali"
  echo " exp/lang2/tri5_ali exp/lang1/nnet3/egs exp/lang2/nnet3/egs"
  echo ""
  echo "Main options (for others, see top of script file)"
  echo "  --config <config-file>                           # config file containing options"
  echo "  --num-jobs <nj>                                  # The maximum number of jobs you want to run in"
  echo "                                                   # parallel (increase this only if you have good disk and"
  echo "                                                   # network speed).  default=6"
  echo "  --cmd (utils/run.pl;utils/queue.pl <queue opts>) # how to run jobs."
  echo "  --samples-per-iter <#samples;400000>             # Target number of egs per archive (option is badly named)"
  echo "  --frames-per-eg <frames;8>                       # number of frames per eg on disk"
  echo "  --left-context <width;4>                         # Number of frames on left side to append for feature input"
  echo "  --right-context <width;4>                        # Number of frames on right side to append for feature input"
  echo "  --num-frames-diagnostic <#frames;4000>           # Number of frames used in computing (train,valid) diagnostics"
  echo "  --num-valid-frames-combine <#frames;10000>       # Number of frames used in getting combination weights at the"
  echo "                                                   # very end."
  echo "  --stage <stage|0>                                # Used to run a partially-completed training process from somewhere in"
  echo "                                                   # the middle."

  exit 1;
fi

num_lang=$1
shift
args=("$@")

echo "$num_lang"
echo "${args[@]}"

if [ ${#args[@]} != $[$num_lang*4] ]; then
    echo "$0: num of input dirs provided for all langs is not compatible with num-langs in input." && exit 1;
fi

# read input data, ali and egs dir per lang
for l in `seq 0 $[$num_lang-1]`; do
    multi_data_dirs[$l]=${args[$l]}
    multi_ali_dirs[$l]=${args[$l+$num_lang]}
    multi_egs_dirs[$l]=${args[$l+2*$num_lang]}
    num_targets_list[$l]=${args[$l+3*$num_lang]}
done

echo "$0: num_targets_list= (${num_targets_list[@]})"
echo "$0: Generate separate egs directory per language for multilingual training."
online_multi_ivector_dirs=(${online_multi_ivector_dirs[@]})

for lang_index in `seq 0 $[$num_lang-1]`; do
    data=${multi_data_dirs[$lang_index]}
    ali_dir=${multi_ali_dirs[$lang_index]}
    egs_dir=${multi_egs_dirs[$lang_index]}
    online_ivector_dir=
    num_targets=${num_targets_list[$lang_index]}
    
    if [ ! -z "${online_multi_ivector_dirs[$lang_index]}" ]; then
        online_ivector_dir=${online_multi_ivector_dirs[$lang_index]}
    fi
    echo online_ivector_dir = $online_ivector_dir

    if [ ! -d "$egs_dir" ]; then

        echo "$0: Generate egs for $ali_dir"
        
        extra_opts=()
        [ ! -z "$cmvn_opts" ] && extra_opts+=(--cmvn-opts "$cmvn_opts")
        [ ! -z "$online_ivector_dir" ] && extra_opts+=(--online-ivector-dir $online_ivector_dir)
        extra_opts+=(--left-context $left_context)
        extra_opts+=(--right-context $right_context)
        
        echo "$0: calling get_egs.sh"
        
        steps/nnet3/get_egs.sh \
            $egs_opts "${extra_opts[@]}" \
            --samples-per-iter $samples_per_iter \
            --stage $stage \
            --cmd "$cmd" \
            $egs_opts \
            --generate-egs-scp true \
            $data \
            $ali_dir \
            $egs_dir \
            || exit 1;
        
    fi
done

