#!/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
# run step 0 to step 5 by default
stage=0
stop_stage=5

# 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
#     - small
#     - medium
#     - large
#  
#     You can download them from
#      - https://dl.fbaipublicfiles.com/librilight/data/small.tar
#      - https://dl.fbaipublicfiles.com/librilight/data/medium.tar
#      - https://dl.fbaipublicfiles.com/librilight/data/large.tar

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 "Running prepare.sh"

log "dl_dir: $dl_dir"

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

if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then
  log "Stage 2: Preprocess Libri-Light manifest"
  mkdir -p data/kmeans
  if [ ! -f data/kmeans/.preprocess_complete ]; then
    python3 ./local/preprocess_librilight.py
    touch data/fbank/.preprocess_complete
  fi
fi

if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then
  log "Stage 3: Split medium and large subset into pieces"
  num_per_split=200000
  split_dir=data/kmeans/medium_split
  if [ ! -f $split_dir/.split_completed ]; then
    lhotse split-lazy ./data/kmeans/librilight_cuts_medium_raw.jsonl.gz $split_dir $num_per_split
    touch $split_dir/.split_completed
  fi
  split_dir=data/kmeans/large_split
  if [ ! -f $split_dir/.split_completed ]; then
    lhotse split-lazy ./data/kmeans/librilight_cuts_large_raw.jsonl.gz $split_dir $num_per_split
    touch $split_dir/.split_completed
  fi
fi

if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then
  log "Stage 4: Extract SSL target for librilight"
  mkdir -p data/fbank
  if [ ! -e data/fbank/.librispeech.done ]; then
    ./local/compute_fbank_librispeech.py
    touch data/fbank/.librispeech.done
  fi
fi