#!/usr/bin/env bash

# fix segmentation fault reported in https://github.com/k2-fsa/icefall/issues/674
export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python

set -eou pipefail

nj=15
stage=0
stop_stage=100

. shared/parse_options.sh || exit 1

# All files generated by this script are saved in "data".
# You can safely remove "data" and rerun this script to regenerate it.
mkdir -p data

log() {
  # This function is from espnet
  local fname=${BASH_SOURCE[1]##*/}
  echo -e "$(date '+%Y-%m-%d %H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
}

if [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then
  log "Stage 0: Prepare gigaspeech dataset."
  mkdir -p data/fbank
  if [ ! -e data/fbank/.gigaspeech.done ]; then
    pushd ../ASR
    ./prepare.sh --stage 0 --stop-stage 9
    ./prepare.sh --stage 11 --stop-stage 11
    popd
    pushd data/fbank
    ln -svf $(realpath ../ASR/data/fbank/gigaspeech_cuts_DEV.jsonl.gz) .
    ln -svf $(realpath ../ASR/data/fbank/gigaspeech_feats_DEV.lca) .
    ln -svf $(realpath ../ASR/data/fbank/gigaspeech_cuts_TEST.jsonl.gz) .
    ln -svf $(realpath ../ASR/data/fbank/gigaspeech_feats_TEST.lca) .
    ln -svf $(realpath ../ASR/data/fbank/gigaspeech_cuts_L.jsonl.gz) .
    ln -svf $(realpath ../ASR/data/fbank/gigaspeech_feats_L.lca) .
    ln -svf $(realpath ../ASR/data/fbank/gigaspeech_cuts_M.jsonl.gz) .
    ln -svf $(realpath ../ASR/data/fbank/gigaspeech_feats_M.lca) .
    ln -svf $(realpath ../ASR/data/fbank/gigaspeech_cuts_S.jsonl.gz) .
    ln -svf $(realpath ../ASR/data/fbank/gigaspeech_feats_S.lca) .
    ln -svf $(realpath ../ASR/data/fbank/gigaspeech_cuts_XS.jsonl.gz) .
    ln -svf $(realpath ../ASR/data/fbank/gigaspeech_feats_XS.lca) .
    ln -svf $(realpath ../ASR/data/fbank/XL_split) .
    ln -svf $(realpath ../ASR/data/fbank/musan_cuts.jsonl.gz) .
    ln -svf $(realpath ../ASR/data/fbank/musan_feats) .
    popd
    pushd data
    ln -svf $(realpath ../ASR/data/lang_bpe_500) .
    popd
    touch data/fbank/.gigaspeech.done
  else
    log "Gigaspeech dataset already exists, skipping."
  fi
fi

if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then
  log "Stage 1: Prepare open commands dataset."
  mkdir -p data/fbank
  if [ ! -e data/fbank/.fluent_speech_commands.done ]; then
    pushd data
    git clone https://github.com/pkufool/open-commands.git
    ln -svf $(realpath ./open-commands/EN/small/commands.txt) commands_small.txt
    ln -svf $(realpath ./open-commands/EN/large/commands.txt) commands_large.txt
    pushd open-commands
    ./script/prepare.sh --stage 2 --stop-stage 2
    ./script/prepare.sh --stage 6 --stop-stage 6
    popd
    popd
    pushd data/fbank
    ln -svf $(realpath ../open-commands/data/fbank/fluent_speech_commands_cuts_large.jsonl.gz) .
    ln -svf $(realpath ../open-commands/data/fbank/fluent_speech_commands_feats_large) .
    ln -svf $(realpath ../open-commands/data/fbank/fluent_speech_commands_cuts_small.jsonl.gz) .
    ln -svf $(realpath ../open-commands/data/fbank/fluent_speech_commands_feats_small) .
    ln -svf $(realpath ../open-commands/data/fbank/fluent_speech_commands_cuts_valid.jsonl.gz) .
    ln -svf $(realpath ../open-commands/data/fbank/fluent_speech_commands_feats_valid) .
    ln -svf $(realpath ../open-commands/data/fbank/fluent_speech_commands_cuts_train.jsonl.gz) .
    ln -svf $(realpath ../open-commands/data/fbank/fluent_speech_commands_feats_train) .
    popd
    touch data/fbank/.fluent_speech_commands.done
  else
    log "Fluent speech commands dataset already exists, skipping."
  fi
fi