#!/usr/bin/env bash

set -eou pipefail

stage=3
stop_stage=3

# We assume dl_dir (download dir) contains the following
# directories and files. If not, they will be downloaded
# by this script automatically.
#
#  - $dl_dir/SPEECHIO_ASR_ZH00000
#      This directory contains the following files downloaded from
#       https://github.com/SpeechColab/Leaderboard
#
#     - metadata.tsv
#     - wav
#     - wav.scp
#     - trans.txt
#

dl_dir=$PWD/download

. 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]}) $*"
}

log "dl_dir: $dl_dir"

if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then
  log "Stage 1: Prepare speechio manifest"
  # We assume that you have downloaded the speechio dataset
  # to $dl_dir
  mkdir -p data/manifests
  if [ ! -e data/manifests/.speechio.done ]; then
    lhotse prepare speechio $dl_dir data/manifests
    touch data/manifests/.speechio.done
  fi
fi

whisper_mel_bins=80
if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then
  log "Stage 2: Compute whisper fbank for speechio"
  if [ ! -f data/fbank/.speechio.done ]; then
    mkdir -p data/fbank
    ./local/compute_fbank_speechio.py --num-mel-bins ${whisper_mel_bins} --whisper-fbank true
    touch data/fbank/.speechio.done
  fi
fi

if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then
  log "Stage 3: Compute kaldi fbank for speechio"
  if [ ! -f data/fbank/.speechio.kaldi.done ]; then
    fbank_dir=data/fbank_kaldi
    mkdir -p $fbank_dir
    ./local/compute_fbank_speechio.py --fbank-dir $fbank_dir
    touch data/fbank/.speechio.kaldi.done
  fi
fi