#!/usr/bin/env bash
set -e

# bash xgiga_tokenize.sh /opt/tiger/sumtest/datas/xgiga en en_XX MSPM test
# bash xgiga_tokenize.sh /opt/tiger/sumtest/datas/xgiga fr fr_XX MSPM
# bash xgiga_tokenize.sh /opt/tiger/sumtest/datas/xgiga zh zh_CN MSPM

DATADIR="$1"
LG="$2"
TAG="$3"
MODE="$4"
DIVISIONS="$5"

if [ -z "$DATADIR" ]; then
    echo "Lose the raw data dir!"
    exit
fi
if [ -z "$LG" ]; then
    echo "Lose the language!"
    exit
fi
if [[ -n "$DIVISIONS" ]] && [[ "$DIVISIONS" == "test" ]]; then
    PARTS=(test)
else
    PARTS=(train dev test)
fi
echo "Parts are : ${PARTS[*]}"

VOCAB_SIZE=32000
DATA="$DATADIR"
TOKEN="$DATADIR/$MODE/$LG"
if [ ! -d "$TOKEN"  ] ; then
    mkdir -p "$TOKEN"
fi

if [[ "$MODE" == "SPM" ]]; then
  if [ ! -f "$DATA/sentence.bpe.$VOCAB_SIZE.$LG.model" ] ; then
    echo "SPM training for Language..."
    if [[ "$LG" == "zh" ]] || [[ "$LG" == "ja" ]] ; then
      python fairseq/scripts/spm_train.py \
      --input="$DATA/lines.txt" \
      --model_prefix="$DATA/sentence.bpe.$VOCAB_SIZE.$LG" \
      --vocab_size=$VOCAB_SIZE \
      --character_coverage=0.995 \
      --model_type=bpe \
      --user_defined_symbols='<q>'
    else
      python fairseq/scripts/spm_train.py \
      --input="$DATA/lines.txt" \
      --model_prefix="$DATA/sentence.bpe.$VOCAB_SIZE.$LG" \
      --vocab_size=$VOCAB_SIZE \
      --character_coverage=1 \
      --model_type=bpe \
      --user_defined_symbols='<q>'
    fi
  fi

  echo "SPM encoding for dataset..."
  for SPLIT in ${PARTS[*]}
    do
      echo "  encoding $DATA/$SPLIT.$LG.doc to $TOKEN/$SPLIT.$LG.spm.doc ..."
      python fairseq/scripts/spm_encode.py --model="$DATA/sentence.bpe.$VOCAB_SIZE.$LG.model" < "$DATA/$SPLIT.$LG.doc" | sed "s/^/<s> /" -e "s/$/ <\/s>/" -e "s/^/[${TAG}] /" > "$TOKEN/$SPLIT.$LG.spm.doc"
      echo "  encoding $DATA/$SPLIT.$LG.sum to $TOKEN/$SPLIT.$LG.spm.sum ..."
      python fairseq/scripts/spm_encode.py --model="$DATA/sentence.bpe.$VOCAB_SIZE.$LG.model" < "$DATA/$SPLIT.$LG.sum" | sed "s/^/<s> /" -e "s/$/ <\/s>/" -e "s/^/[${TAG}] /"> "$TOKEN/$SPLIT.$LG.spm.sum"
    done

  echo "Changing vocab to dict..."
  python -m preprocess.vocab2dict -i "$DATA/sentence.bpe.$VOCAB_SIZE.$LG.vocab" -o "$DATA/$LG.dict.txt" -l "$LG"
  DICT="$DATA/$LG.dict.txt"


elif [[ "$MODE" == "MSPM" ]]; then
  # MBART=/data00/home/wangdanqing.122/workshop/MultiLingual/train/pretrain/mbart/mbart.cc25
  MBART=/home/tiger/mbart.cc25.v2
  MODEL=$MBART/sentence.bpe.model
  # DICT=$MBART/dict.txt
  DICT=$MBART/dict.txt


  echo "MSPM encoding for dataset..."
  for SPLIT in ${PARTS[*]}
    do
      echo "  encoding $DATA/$SPLIT.x.$LG to $DATA/$SPLIT.x.$LG.spm ..."
      python3 fairseq/scripts/spm_encode.py --model=$MODEL < "$DATA/$SPLIT.x.$LG" | sed -e "s/< q >/ <\/s>/g" > $DATA/$SPLIT.x.$LG.spm

      echo "  encoding $DATA/$SPLIT.y.$LG to $DATA/$SPLIT.y.$LG.spm ..."
      python3 fairseq/scripts/spm_encode.py --model=$MODEL < "$DATA/$SPLIT.y.$LG" | sed -e "s/< q >/ <\/s>/g" > $DATA/$SPLIT.y.$LG.spm
    done


elif [[ "$MODE" == "BPE" ]]; then
  wget -N 'https://dl.fbaipublicfiles.com/fairseq/gpt2_bpe/encoder.json'
  wget -N 'https://dl.fbaipublicfiles.com/fairseq/gpt2_bpe/vocab.bpe'
  wget -N 'https://dl.fbaipublicfiles.com/fairseq/gpt2_bpe/dict.txt'

  DICT="dict.txt"
  for SPLIT in ${PARTS[*]}
  do
    for TYPE in doc sum
    do
      python -m examples.roberta.multiprocessing_bpe_encoder \
        --encoder-json encoder.json \
        --vocab-bpe vocab.bpe \
        --inputs "$DATA/$SPLIT.$LG.$TYPE" \
        --outputs "$TOKEN/$SPLIT.$LG.bpe.$TYPE" \
        --workers 60 \
        --keep-empty
    done
  done

elif [[ "$MODE" == "PMNMT" ]]; then
  # MODEL=toutiaomt/codes.bpe.32000
  # DICT=toutiaomt/vocab.bpe.32000
  LANGID="LANG_TOK_"$(echo "${LG}" | tr '[a-z]' '[A-Z]')
  MODEL=~/mCOLT/vocab/codes.bpe.32000
  DICT=~/mCOLT/vocab/dict.txt

  echo "PMNMT BPE encoding for dataset..."
  for SPLIT in ${PARTS[*]}
    do
      echo "  encoding $DATA/$SPLIT.$LG.doc to $TOKEN/$SPLIT.$LG.bpe.doc ..."
      python toutiaomt/text_processing/bpe.py --codes=$MODEL --input $DATA/$SPLIT.$LG.doc --output $TOKEN/$SPLIT.$LG.nonlang.bpe.doc --threads 32
      sed -e 's/^/'${LANGID}' /' -e "s/@@ <@@ q@@ >@@ /@@ <\/s> /g" -e "s/“@@/\"@@/g" -e "s/”@@/\"@@/g" $TOKEN/$SPLIT.$LG.nonlang.bpe.doc > $TOKEN/$SPLIT.$LG.bpe.doc
      echo "  encoding $DATA/$SPLIT.$LG.sum to $TOKEN/$SPLIT.$LG.bpe.sum ..."
      python toutiaomt/text_processing/bpe.py --codes=$MODEL --input $DATA/$SPLIT.$LG.sum --output $TOKEN/$SPLIT.$LG.nonlang.bpe.sum --threads 32
      sed -e 's/^/'${LANGID}' /' -e "s/@@ <@@ q@@ >@@ /@@ <\/s> /g" -e "s/“@@/\"@@/g" -e "s/”@@/\"@@/g" $TOKEN/$SPLIT.$LG.nonlang.bpe.sum > $TOKEN/$SPLIT.$LG.bpe.sum
    done


fi
