#!/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=-1
stop_stage=100

# We assume dl_dir (download dir) contains the following
# directories and files. If not, they will be downloaded
# by this script automatically.
#
#  - $dl_dir/librilight
#      You can find small, medium, large, etc. inside it.
#
#  - $dl_dir/libriheavy
#      You can find libriheavy_cuts_small.jsonl.gz, libriheavy_cuts_medium.jsonl.gz, etc. inside it.
dl_dir=$PWD/download

. shared/parse_options.sh || exit 1

# vocab size for sentence piece models.
# It will generate data/lang_bpe_xxx,
# data/lang_bpe_yyy if the array contains xxx, yyy
vocab_sizes=(
  4000
)

# 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
tokens_dir=data/tokens
manifests_dir=data/manifests

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: Download audio data."
  # If you have pre-downloaded it to /path/to/librilight,
  # you can create a symlink
  #
  #   ln -sfv /path/to/librilight $dl_dir/librilight
  #
  mkdir -p $dl_dir/librilight
  for subset in small medium large; do
    log "Downloading ${subset} subset."
    if [ ! -d $dl_dir/librilight/${subset} ]; then
      wget -P $dl_dir/librilight -c https://dl.fbaipublicfiles.com/librilight/data/${subset}.tar 
      tar xf $dl_dir/librilight/${subset}.tar -C $dl_dir/librilight
    else
      log "Skipping download, ${subset} subset exists."
    fi
  done
fi

if [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then
  log "Stage 0: Download manifests from huggingface."

  # If you have pre-downloaded it to /path/to/libriheavy,
  # you can create a symlink
  #
  #   ln -sfv /path/to/libriheavy $dl_dir/libriheavy
  #
  mkdir -p $dl_dir/libriheavy
  for subset in small medium large dev test_clean test_other; do
    if [ ! -e $dl_dir/libriheavy/libriheavy_cuts_${subset}.jsonl.gz ]; then
      log "Downloading ${subset} subset."
      wget -P $dl_dir/libriheavy -c https://huggingface.co/datasets/pkufool/libriheavy/resolve/main/libriheavy_cuts_${subset}.jsonl.gz
    else
      log "Skipping download, ${subset} subset exists."
    fi
  done
fi

if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then
  log "Stage 1: Prepare Libriheavy manifests"
  mkdir -p $manifests_dir
  for subset in small medium large dev test_clean test_other; do
    if [ ! -e $manifests_dir/libriheavy_cuts_${subset}.jsonl.gz ]; then
      log "Prepare manifest for subset : ${subset}"
      ./local/prepare_manifest.py $dl_dir/libriheavy/libriheavy_cuts_${subset}.jsonl.gz $manifests_dir
    fi
  done
fi

num_per_split=200000
if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then
  log "Stage 2: Split medium and large subsets."
  for subset in medium large; do
    log "Spliting subset : $subset"
    split_dir=$manifests_dir/libriheavy_${subset}_split
    mkdir -p $split_dir
    if [ ! -e $split_dir/.split_completed ]; then
      lhotse split-lazy $manifests_dir/libriheavy_cuts_${subset}.jsonl.gz $split_dir $num_per_split
      touch $split_dir/.split_completed
    fi
  done
fi

if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then
  log "Stage 3: Train BPE model for normalized text"

  if [ ! -f data/texts ]; then
    gunzip -c $manifests_dir/libriheavy_cuts_medium.jsonl.gz \
      | jq '.supervisions[].text' | sed 's/"//;s/\\//g;s/"$//' \
      | ./local/norm_text.py > data/texts
  fi

  for vocab_size in ${vocab_sizes[@]}; do
    lang_dir=data/lang_bpe_${vocab_size}
    mkdir -p $lang_dir

    cp data/texts $lang_dir/text

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

if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then
  log "Stage 4: Extract speech tokens."
  for subset in small medium large; do
    log "Extract speech tokens for subset: $subset"
    output_dir=$tokens_dir/libriheavy_${subset}
    mkdir -p $tokens_dir
    if [ ! -e $tokens_dir/.extract_completed ]; then
      torchrun --nproc_per_node=8 \
        --nnodes=1 \
        --rdzv_id=2024 \
        --rdzv_backend="c10d" \
        --rdzv_endpoint="localhost:0" \
        `which s3tokenizer` \
          --cuts_path $manifests_dir/libriheavy_cuts_${subset}.jsonl.gz \
          --device "cuda" \
          --output_dir $output_dir \
          --batch_size 32 \
          --model "speech_tokenizer_v1"
      cat $output_dir/part* | gzip > $output_dir/libriheavy_${subset}.jsonl.gz && rm -rf $output_dir
      touch $output_dir/..extract_completed
    fi
  done
fi

if [ $stage -le 5 ] && [ $stop_stage -ge 5 ]; then
  log "Stage 5: Attach speech tokens."
  for subset in small medium large; do
    log "Attach speech tokens for subset: $subset"
    if [ ! -e $tokens_dir/libriheavy_cuts_${subset}.jsonl.gz ]; then
      ./local/attach_speech_tokens.py --subset $subset
    fi
  done
fi