#!/usr/bin/env bash

set -eou pipefail

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/GRID
#      You can find lip, audio, align_text inside it.
#
#  - $dl_dir/lm
#      This directory contains the language model(LM) downloaded from
#      https://huggingface.co/luomingshuang/grid_lm.
#      About how to get these LM files, you can know it 
#      from https://github.com/luomingshuang/Train_LM_with_kaldilm.
#	
#	    - lm_3_gram.arpa
#     - lm_4_gram.arpa
#

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: Download LM"
  # We assume that you have installed the git-lfs, if not, you could install it
  # using: `sudo apt-get install git-lfs && git-lfs install`
  #[ ! -e $dl_dir/lm ] && mkdir -p $dl_dir/lm
  #git clone https://huggingface.co/luomingshuang/grid_lm $dl_dir/lm
  #cd $dl_dir/lm && git lfs pull

  # You can also use the following commands to download the lm files
  wget -P $dl_dir/lm https://huggingface.co/luomingshuang/grid_lm/resolve/main/lm_3_gram.arpa
  wget -P $dl_dir/lm https://huggingface.co/luomingshuang/grid_lm/resolve/main/lm_4_gram.arpa

  # Because the texts among the samples in GRID are very similar,
  # the lm_4_gram.arpa is nearly no use for decoding when use LM.
fi

if [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then
  log "Stage 0: Download data"
  # The process of extracting lip region takes much time.
  # Here, we provide the processed data (lip region) for using. 
  # So you can run this recipe quickly and easily.
  #
  # If you want to know more details about getting lip region,
  # You can have a look at https://github.com/Fengdalu/LipNet-PyTorch/tree/master/scripts

  [ ! -e $dl_dir/GRID ] && mkdir -p $dl_dir/GRID

  # Download the GRID lip region data and text
  # You can use the following commands to download the processed lip region data and text
  wget -P $dl_dir/GRID https://huggingface.co/datasets/luomingshuang/grid_lip_160_80/resolve/main/GRID_LIP_160x80_TXT.zip.00
  wget -P $dl_dir/GRID https://huggingface.co/datasets/luomingshuang/grid_lip_160_80/resolve/main/GRID_LIP_160x80_TXT.zip.01
  wget -P $dl_dir/GRID https://huggingface.co/datasets/luomingshuang/grid_lip_160_80/resolve/main/GRID_LIP_160x80_TXT.zip.02
  wget -P $dl_dir/GRID https://huggingface.co/datasets/luomingshuang/grid_lip_160_80/resolve/main/GRID_LIP_160x80_TXT.zip.03
  wget -P $dl_dir/GRID https://huggingface.co/datasets/luomingshuang/grid_lip_160_80/resolve/main/GRID_LIP_160x80_TXT.zip.04

  cat $dl_dir/GRID/GRID_LIP_160x80_TXT.zip.* > $dl_dir/GRID/GRID_LIP_160x80_TXT.zip
  unzip $dl_dir/GRID/GRID_LIP_160x80_TXT.zip -d $dl_dir/GRID/
  rm -rf $dl_dir/GRID/GRID_LIP_160x80_TXT.zip
  
  # Download the GRID audio data
  wget -P $dl_dir/GRID https://huggingface.co/datasets/luomingshuang/GRID_audio/resolve/main/audio_25k.zip
  unzip $dl_dir/GRID/audio_25k.zip -d $dl_dir/GRID/
  rm -rf $dl_dir/GRID/audio_25k.zip

  # Download the spliting files for train and val
  # Here, we just consider the unseen case, which means 
  # that there is no common speakers among train and val.
  wget -P $dl_dir/GRID https://huggingface.co/datasets/luomingshuang/GRID_text/resolve/main/unseen_train.txt
  wget -P $dl_dir/GRID https://huggingface.co/datasets/luomingshuang/GRID_text/resolve/main/unseen_val.txt
fi

if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then
  log "Stage 2: Prepare character-based lang"
  lang_dir=data/lang_character
  mkdir -p $lang_dir

  ./local/prepare_lexicon.py \
   --samples-txt $dl_dir/GRID/unseen_train.txt \
   --align-dir $dl_dir/GRID/GRID_align_txt \
   --lang-dir $lang_dir

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

if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then
  log "Stage 3: Prepare G"
  # We assume you have installed kaldilm, if not, please install
  # it using: pip install kaldilm

  mkdir -p data/lm
  if [ ! -f data/lm/G_3_gram.fst.txt ]; then
    # It is used in building HLG
    python3 -m kaldilm \
      --read-symbol-table="data/lang_character/words.txt" \
      --disambig-symbol='#0' \
      --max-order=3 \
      $dl_dir/lm/lm_3_gram.arpa > data/lm/G_3_gram.fst.txt
  fi

  if [ ! -f data/lm/G_4_gram.fst.txt ]; then
    # It is used for LM rescoring
    python3 -m kaldilm \
      --read-symbol-table="data/lang_character/words.txt" \
      --disambig-symbol='#0' \
      --max-order=4 \
      $dl_dir/lm/lm_4_gram.arpa > data/lm/G_4_gram.fst.txt
  fi
fi

if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then
  log "Stage 4: Compile HLG"
  ./local/compile_hlg.py --lang-dir data/lang_character
fi