#!/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 wenetspeech dataset."
  mkdir -p data/fbank
  if [ ! -e data/fbank/.wenetspeech.done ]; then
    pushd ../ASR
    ./prepare.sh --stage 0 --stop-stage 17
    ./prepare.sh --stage 22 --stop-stage 22
    popd
    pushd data/fbank
    ln -svf $(realpath ../ASR/data/fbank/cuts_DEV.jsonl.gz) .
    ln -svf $(realpath ../ASR/data/fbank/feats_DEV.lca) .
    ln -svf $(realpath ../ASR/data/fbank/cuts_TEST_NET.jsonl.gz) .
    ln -svf $(realpath ../ASR/data/fbank/feats_TEST_NET.lca) .
    ln -svf $(realpath ../ASR/data/fbank/cuts_TEST_MEETING.jsonl.gz) .
    ln -svf $(realpath ../ASR/data/fbank/feats_TEST_MEETING.lca) .
    ln -svf $(realpath ../ASR/data/fbank/cuts_L.jsonl.gz) .
    ln -svf $(realpath ../ASR/data/fbank/L_split_1000) .
    ln -svf $(realpath ../ASR/data/fbank/cuts_M.jsonl.gz) .
    ln -svf $(realpath ../ASR/data/fbank/M_split_1000) .
    ln -svf $(realpath ../ASR/data/fbank/cuts_S.jsonl.gz) .
    ln -svf $(realpath ../ASR/data/fbank/S_split_1000) .
    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_partial_tone) .
    popd
    touch data/fbank/.wenetspeech.done
  else
    log "WenetSpeech 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/.cn_speech_commands.done ]; then
    pushd data
    git clone https://github.com/pkufool/open-commands.git
    ln -svf $(realpath ./open-commands/CN/small/commands.txt) commands_small.txt
    ln -svf $(realpath ./open-commands/CN/large/commands.txt) commands_large.txt
    pushd open-commands
    ./scripts/prepare.sh --stage 1 --stop-stage 1
    ./scripts/prepare.sh --stage 3 --stop-stage 5
    popd
    popd
    pushd data/fbank
    ln -svf $(realpath ../open-commands/data/fbank/cn_speech_commands_cuts_large.jsonl.gz) .
    ln -svf $(realpath ../open-commands/data/fbank/cn_speech_commands_feats_large) .
    ln -svf $(realpath ../open-commands/data/fbank/cn_speech_commands_cuts_small.jsonl.gz) .
    ln -svf $(realpath ../open-commands/data/fbank/cn_speech_commands_feats_small) .
    ln -svf $(realpath ../open-commands/data/fbank/nihaowenwen_cuts_dev.jsonl.gz) .
    ln -svf $(realpath ../open-commands/data/fbank/nihaowenwen_feats_dev) .
    ln -svf $(realpath ../open-commands/data/fbank/nihaowenwen_cuts_test.jsonl.gz) .
    ln -svf $(realpath ../open-commands/data/fbank/nihaowenwen_feats_test) .
    ln -svf $(realpath ../open-commands/data/fbank/nihaowenwen_cuts_train.jsonl.gz) .
    ln -svf $(realpath ../open-commands/data/fbank/nihaowenwen_feats_train) .
    ln -svf $(realpath ../open-commands/data/fbank/xiaoyun_cuts_clean.jsonl.gz) .
    ln -svf $(realpath ../open-commands/data/fbank/xiaoyun_feats_clean.lca) .
    ln -svf $(realpath ../open-commands/data/fbank/xiaoyun_cuts_noisy.jsonl.gz) .
    ln -svf $(realpath ../open-commands/data/fbank/xiaoyun_feats_noisy.lca) .
    popd
    touch data/fbank/.cn_speech_commands.done
  else
    log "CN speech commands dataset already exists, skipping."
  fi
fi