#!/usr/bin/env bash

# adapted from: https://github.com/eole-nlp/eole/blob/357b7d3/recipes/wmt17/prepare_wmt_ende_data.sh

# Set default values for method and encoding
METHOD="bpe"  # Default method
ENCODE=true   # Default is to encode the data

suffix=""

# Parse command line arguments
while [[ "$#" -gt 0 ]]; do
  case $1 in
    --method)
      if [[ "$2" == "bpe" || "$2" == "sentencepiece" ]]; then
        METHOD="$2"
        shift
      else
        echo "Invalid method specified. Use '--method bpe' or '--method sentencepiece'."
        exit 1
      fi
      ;;
    --encode)
      if [[ "$2" == "true" || "$2" == "false" ]]; then
        ENCODE="$2"
        shift
      else
        echo "Invalid value for --encode. Use 'true' or 'false'."
        exit 1
      fi
      ;;
    *)
      echo "Unknown parameter passed: $1"
      exit 1
      ;;
  esac
  shift
done

# Check if subword-nmt is installed for BPE
if [[ "$METHOD" == "bpe" ]] && ! command -v subword-nmt &>/dev/null; then
  echo "Please install Subword NMT: pip3 install subword-nmt"
  exit 2
fi

# Check if sentencepiece is installed for SentencePiece
if [[ "$METHOD" == "sentencepiece" ]] && (! command -v spm_train &>/dev/null || ! command -v spm_encode &>/dev/null); then
  echo "Please install SentencePiece: https://github.com/google/sentencepiece?tab=readme-ov-file#build-and-install-sentencepiece-command-line-tools-from-c-source"
  exit 2
fi

cd toy-lanotes

# prepare data
if true; then
  # Learn subword model based on selected method
  if [[ "$METHOD" == "sentencepiece" ]]; then
    # Using SentencePiece
    cat src-train.txt tgt-train.txt > train.all
    spm_train --input=train.all --model_prefix=spm --vocab_size=32000 --character_coverage=1.0 --model_type=unigram
  else
    # Using BPE with subword-nmt
    cat src-train.txt tgt-train.txt | subword-nmt learn-bpe -s 32000 > codes
  fi

  # Encode data if --encode=true
  if [[ "$ENCODE" == "true" ]]; then
    suffix=".encoded.filter"
    if [[ "$METHOD" == "sentencepiece" ]]; then
      # Apply SentencePiece encoding
      for LANG in src tgt; do
        cat $LANG-train.txt | spm_encode --model=spm.model > train.$LANG.encoded
        for SET in dev test; do
          spm_encode --model=spm.model < $SET.$LANG > $SET.$LANG.encoded
        done
      done
    else
      # Apply BPE encoding with subword-nmt
      for LANG in src tgt; do
        cat $LANG-train.txt | subword-nmt apply-bpe -c codes > train.$LANG.encoded
        #for SET in dev test; do
        #  subword-nmt apply-bpe -c codes < $SET.$LANG > $SET.$LANG.encoded
        #done
      done
    fi

    # Filter the training data
    python3 ../filter_train.py
  else
    echo "Encoding skipped as per --encode=false"
  fi
fi

# Shuffle the training data and create validation set of 5,000 lines
paste -d '\t' train.src$suffix train.tgt$suffix \
| shuf \
| awk -F'\t' -v k=350 '
{
    src[NR] = $1
    tgt[NR] = $2
}
END {
    split_point = NR - k
    for (i = 1; i <= NR; i++) {
        if (i <= split_point) {
            print src[i] > "train.src'$suffix'.shuf"
            print tgt[i] > "train.tgt'$suffix'.shuf"
        } else {
            print src[i] > "valid.src'$suffix'.shuf"
            print tgt[i] > "valid.tgt'$suffix'.shuf"
        }
    }
}'