#!/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=10
stage=0
stop_stage=100

version=v1.0
tgt_lang=de
dl_dir=$PWD/download

must_c_dir=$dl_dir/must-c/$version/en-$tgt_lang/data

# We assume dl_dir (download dir) contains the following
# directories and files.
#  - $dl_dir/must-c/$version/en-$tgt_lang/data/{dev,train,tst-COMMON,tst-HE}
#
# Please go to https://ict.fbk.eu/must-c-releases/
# to download and untar the dataset if you have not already done this.

#  - $dl_dir/musan
#      This directory contains the following directories downloaded from
#       http://www.openslr.org/17/
#
#     - music
#     - noise
#     - speech

. shared/parse_options.sh || exit 1

# vocab size for sentence piece models.
# It will generate
#  data/lang_bpe_${tgt_lang}_xxx
#  data/lang_bpe_${tgt_lang}_yyy
# if the array contains xxx, yyy
vocab_sizes=(
  # 5000
  # 2000
  # 1000
  500
)

# 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 [ ! -d $must_c_dir ]; then
  log "$must_c_dir does not exist"
  exit 1
fi

for d in dev train tst-COMMON tst-HE; do
  if [ ! -d $must_c_dir/$d ]; then
    log "$must_c_dir/$d does not exist!"
    exit 1
  fi
done

if [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then
  log "Stage 0: Download musan"
  if [ ! -d $dl_dir/musan ]; then
    lhotse download musan $dl_dir
  fi
fi

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

if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then
  log "Stage 2: Prepare must-c $version manifest for target language $tgt_lang"
  mkdir -p data/manifests/$version
  if [ ! -e data/manifests/$version/.${tgt_lang}.manifests.done ]; then
    lhotse prepare must-c \
      -j $nj \
      --tgt-lang $tgt_lang \
      $dl_dir/must-c/$version/ \
      data/manifests/$version/

    touch data/manifests/$version/.${tgt_lang}.manifests.done
  fi
fi

if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then
  log "Stage 3: Text normalization for $version with target language $tgt_lang"
  if [ ! -f ./data/manifests/$version/.$tgt_lang.norm.done ]; then
    ./local/preprocess_must_c.py \
      --manifest-dir ./data/manifests/$version/ \
      --tgt-lang $tgt_lang
    touch ./data/manifests/$version/.$tgt_lang.norm.done
  fi
fi

if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then
  log "Stage 4: Compute fbank for musan"
  mkdir -p data/fbank
  if [ ! -e data/fbank/.musan.done ]; then
    ./local/compute_fbank_musan.py
    touch data/fbank/.musan.done
  fi
fi

if [ $stage -le 5 ] && [ $stop_stage -ge 5 ]; then
  log "Stage 5: Compute fbank for $version with target language $tgt_lang"
  mkdir -p data/fbank/$version/
  if [ ! -e data/fbank/$version/.$tgt_lang.done ]; then
    ./local/compute_fbank_must_c.py \
      --in-dir ./data/manifests/$version/ \
      --out-dir ./data/fbank/$version/ \
      --tgt-lang $tgt_lang \
      --num-jobs $nj

    ./local/compute_fbank_must_c.py \
      --in-dir ./data/manifests/$version/ \
      --out-dir ./data/fbank/$version/ \
      --tgt-lang $tgt_lang \
      --num-jobs $nj \
      --perturb-speed 1

    touch data/fbank/$version/.$tgt_lang.done
  fi
fi

if [ $stage -le 6 ] && [ $stop_stage -ge 6 ]; then
  log "Stage 6: Prepare BPE based lang for $version with target language $tgt_lang"

  for vocab_size in ${vocab_sizes[@]}; do
    lang_dir=data/lang_bpe_${vocab_size}/$version/$tgt_lang/
    mkdir -p $lang_dir
    if [ ! -f $lang_dir/transcript_words.txt ]; then
      ./local/get_text.py ./data/fbank/$version/must_c_feats_en-${tgt_lang}_train.jsonl.gz > $lang_dir/transcript_words.txt
    fi

    if [ ! -f $lang_dir/words.txt ]; then
      ./local/get_words.py $lang_dir/transcript_words.txt > $lang_dir/words.txt
    fi

    if [ ! -f $lang_dir/bpe.model ]; then
      ./local/train_bpe_model.py \
        --lang-dir $lang_dir \
        --vocab-size $vocab_size \
        --transcript $lang_dir/transcript_words.txt
    fi

    if [ ! -f $lang_dir/L_disambig.pt ]; then
      ./local/prepare_lang_bpe.py --lang-dir $lang_dir

      log "Validating $lang_dir/lexicon.txt"
      ./local/validate_bpe_lexicon.py \
        --lexicon $lang_dir/lexicon.txt \
        --bpe-model $lang_dir/bpe.model
    fi
  done
fi