#!/usr/bin/env bash

# fix segmentation fault reported in https://github.com/k2-fsa/icefall/issues/674
export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python

set -euxo pipefail

nj=20
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/voxpopuli/raw_audios/$lang/$year
#      This directory contains *.ogg files with audio downloaded and extracted from archives:
#       https://dl.fbaipublicfiles.com/voxpopuli/audios/${lang}_${year}.tar
#
#  - Note: the voxpopuli transcripts are downloaded to a ${tmp} folder
#    as part of `lhotse prepare voxpopuli` from:
#       https://dl.fbaipublicfiles.com/voxpopuli/annotations/asr/asr_${lang}.tsv.gz
#
#  - $dl_dir/musan
#      This directory contains the following directories downloaded from
#       http://www.openslr.org/17/
#
#     - music
#     - noise
#     - speech

dl_dir=$PWD/download
#dl_dir=/mnt/matylda6/szoke/EU-ASR/DATA  # BUT

musan_dir=${dl_dir}/musan
#musan_dir=/mnt/matylda2/data/MUSAN  # BUT

# Choose value from ASR_LANGUAGES:
#
# [ "en", "de", "fr", "es", "pl", "it", "ro", "hu", "cs", "nl", "fi", "hr",
# "sk", "sl", "et", "lt" ]
#
# See ASR_LANGUAGES in:
# https://github.com/lhotse-speech/lhotse/blob/c5f26afd100885b86e4244eeb33ca1986f3fa923/lhotse/recipes/voxpopuli.py#L54C4-L54C4
lang=en

task=asr

. shared/parse_options.sh || exit 1

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

# All files generated by this script are saved in "data/${lang}".
# You can safely remove "data/${lang}" and rerun this script to regenerate it.
mkdir -p data/${lang}

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"
log "musan_dir: $musan_dir"
log "task: $task, lang: $lang"

if [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then
  log "Stage 0: Download data"

  # If you have pre-downloaded it to /path/to/$release,
  # you can create a symlink
  #
  #   ln -sfv /path/to/$release $dl_dir/$release
  #
  if [ ! -d $dl_dir/voxpopuli/raw_audios/${lang} ]; then
    lhotse download voxpopuli --subset $lang $dl_dir/voxpopuli
  fi

  # If you have pre-downloaded it to /path/to/musan,
  # you can create a symlink
  #
  #   ln -sfv /path/to/musan $dl_dir/
  #
  if [ ! -d $musan_dir/musan ]; then
    lhotse download musan $musan_dir
  fi
fi

if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then
  log "Stage 1: Prepare VoxPopuli manifest"
  # We assume that you have downloaded the VoxPopuli corpus
  # to $dl_dir/voxpopuli
  if [ ! -e data/manifests/.voxpopuli-${task}-${lang}.done ]; then
    # Warning : it requires Internet connection (it downloads transcripts to ${tmpdir})
    lhotse prepare voxpopuli --task asr --lang $lang -j $nj $dl_dir/voxpopuli data/manifests
    touch data/manifests/.voxpopuli-${task}-${lang}.done
  fi
fi

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

if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then
  log "Stage 3: Preprocess VoxPopuli manifest"
  mkdir -p data/fbank
  if [ ! -e data/fbank/.voxpopuli-${task}-${lang}-preprocess_complete ]; then
    # recordings + supervisions -> cutset
    ./local/preprocess_voxpopuli.py --task $task --lang $lang \
        --use-original-text True
    touch data/fbank/.voxpopuli-${task}-${lang}-preprocess_complete
  fi
fi

if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then
  log "Stage 4: Compute fbank for dev and test subsets of VoxPopuli"
  mkdir -p data/fbank
  for dataset in "dev" "test"; do
    if [ ! -e data/fbank/.voxpopuli-${task}-${lang}-${dataset}.done ]; then
      ./local/compute_fbank.py --src-dir data/fbank --output-dir data/fbank \
          --num-jobs 50 --num-workers ${nj} \
          --prefix "voxpopuli-${task}-${lang}" \
          --dataset ${dataset} \
          --trim-to-supervisions True
      touch data/fbank/.voxpopuli-${task}-${lang}-${dataset}.done
    fi
  done
fi

if [ $stage -le 5 ] && [ $stop_stage -ge 5 ]; then
  log "Stage 5: Compute fbank for train set of VoxPopuli"
  if [ ! -e data/fbank/.voxpopuli-${task}-${lang}-train.done ]; then
    ./local/compute_fbank.py --src-dir data/fbank --output-dir data/fbank \
        --num-jobs 100 --num-workers ${nj} \
        --prefix "voxpopuli-${task}-${lang}" \
        --dataset train \
        --trim-to-supervisions True \
        --speed-perturb True
    touch data/fbank/.voxpopuli-${task}-${lang}-train.done
  fi
fi

if [ $stage -le 6 ] && [ $stop_stage -ge 6 ]; then
  log "Stage 6: Validate fbank manifests for VoxPopuli"
  for dataset in "dev" "test" "train"; do
    mkdir -p data/fbank/log/
    ./local/validate_cutset_manifest.py \
      data/fbank/voxpopuli-asr-en_cuts_${dataset}.jsonl.gz \
      2>&1 | tee data/fbank/log/validate_voxpopuli-asr-en_cuts_${dataset}.log
  done
fi

if [ $stage -le 7 ] && [ $stop_stage -ge 7 ]; then
  log "Stage 7: 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 8 ] && [ $stop_stage -ge 8 ]; then
  log "Stage 8: Prepare BPE based lang"

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

    if [ ! -f $lang_dir/transcript_words.txt ]; then
      log "Generate data for BPE training"
      file=$(
        find "data/fbank/voxpopuli-${task}-${lang}_cuts_train.jsonl.gz"
      )
      local/text_from_manifest.py $file >$lang_dir/transcript_words.txt
      # gunzip -c ${file} | awk -F '"' '{print $30}' > $lang_dir/transcript_words.txt

      # Ensure space only appears once
      #sed -i 's/\t/ /g' $lang_dir/transcript_words.txt
      #sed -i 's/[ ][ ]*/ /g' $lang_dir/transcript_words.txt
    fi

    if [ ! -f $lang_dir/words.txt ]; then
      cat $lang_dir/transcript_words.txt | sed 's/ /\n/g' \
        | sort -u | sed '/^$/d' > $lang_dir/words.txt
      (echo '!SIL'; echo '<SPOKEN_NOISE>'; echo '<UNK>'; ) |
        cat - $lang_dir/words.txt | sort | uniq | awk '
        BEGIN {
          print "<eps> 0";
        }
        {
          if ($1 == "<s>") {
            print "<s> is in the vocabulary!" | "cat 1>&2"
            exit 1;
          }
          if ($1 == "</s>") {
            print "</s> is in the vocabulary!" | "cat 1>&2"
            exit 1;
          }
          printf("%s %d\n", $1, NR);
        }
        END {
          printf("#0 %d\n", NR+1);
          printf("<s> %d\n", NR+2);
          printf("</s> %d\n", NR+3);
        }' > $lang_dir/words || exit 1;
      mv $lang_dir/words $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

    if [ ! -f $lang_dir/L.fst ]; then
      log "Converting L.pt to L.fst"
      ./shared/convert-k2-to-openfst.py \
        --olabels aux_labels \
        $lang_dir/L.pt \
        $lang_dir/L.fst
    fi

    if [ ! -f $lang_dir/L_disambig.fst ]; then
      log "Converting L_disambig.pt to L_disambig.fst"
      ./shared/convert-k2-to-openfst.py \
        --olabels aux_labels \
        $lang_dir/L_disambig.pt \
        $lang_dir/L_disambig.fst
    fi
  done
fi