# Copyright (c) 2019-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#

# usage: ./get_data.sh --src <src language> --tgt <tgt language>
# data should be in ./data/<src language>-<tgt language>
set -e


#
# Data preprocessing configuration
#
N_MONO=130000  # number of monolingual sentences for each language
CODES=5000
N_THREADS=8    # number of threads in data preprocessing


#
# Read arguments
#
POSITIONAL=()
while [[ $# -gt 0 ]]
do
key="$1"
case $key in
  --src)
    SRC="$2"; shift 2;;
  --tgt)
    TGT="$2"; shift 2;;
  --reload_codes)
    RELOAD_CODES="$2"; shift 2;;
  --reload_vocab)
    RELOAD_VOCAB="$2"; shift 2;;
  --codes)
    CODES="$2"; shift 2;;
  --data)
    DATA_PATH="$2"; shift 2;;
  --n_threads)
    N_THREADS="$2"; shift 2;;
  *)
  POSITIONAL+=("$1")
  shift
  ;;
esac
done
set -- "${POSITIONAL[@]}"


#
# Check parameters
#
if [ "$SRC" == "" ]
then
  echo "--src not provided"
  exit
fi
if [ "$TGT" == "" ]
then
  echo "--tgt not provided"
  exit
fi
if [ "$SRC" == "$TGT" ]
then
  echo "source and target cannot be identical"
  exit
fi
if [ "$SRC" \> "$TGT" ]
then
	echo "please ensure SRC < TGT"
	exit
fi
if [ "$RELOAD_CODES" != "" ] && [ ! -f "$RELOAD_CODES" ]
then
  echo "cannot locate BPE codes"
  exit
fi
if [ "$RELOAD_VOCAB" != "" ] && [ ! -f "$RELOAD_VOCAB" ]
then
  echo "cannot locate vocabulary"
  exit
fi
if [ "$RELOAD_CODES" == "" -a "$RELOAD_VOCAB" != "" -o "$RELOAD_CODES" != "" -a "$RELOAD_VOCAB" == "" ]
then
  echo "BPE codes should be provided if and only if vocabulary is also provided"
  exit
fi

# main paths
MAIN_PATH=$PWD
TOOLS_PATH=${MAIN_PATH}/XLM/tools
DATA_PATH_FOR=$DATA_PATH/${SRC}-${TGT}
DATA_PATH_REV=$DATA_PATH/${TGT}-${SRC}
if [ ! -d "$DATA_PATH_FOR" ]; then
	DATA_PATH=$DATA_PATH_REV
	echo "New data path: $DATA_PATH"
	if [ ! -d "$DATA_PATH" ]; then
		echo "Neither $DATA_PATH_FOR nor $DATA_PATH_REV do not exist"
		exit 1
	fi
else
	DATA_PATH=$DATA_PATH_FOR
fi

PROC_PATH=$DATA_PATH/processed

# create paths
mkdir -p $TOOLS_PATH
mkdir -p $DATA_PATH
mkdir -p $PROC_PATH

# moses
MOSES=$TOOLS_PATH/mosesdecoder
REPLACE_UNICODE_PUNCT=$MOSES/scripts/tokenizer/replace-unicode-punctuation.perl
NORM_PUNC=$MOSES/scripts/tokenizer/normalize-punctuation.perl
REM_NON_PRINT_CHAR=$MOSES/scripts/tokenizer/remove-non-printing-char.perl
TOKENIZER=$MOSES/scripts/tokenizer/tokenizer.perl
INPUT_FROM_SGM=$MOSES/scripts/ems/support/input-from-sgm.perl

# fastBPE
FASTBPE_DIR=$TOOLS_PATH/fastBPE
FASTBPE=$FASTBPE_DIR/fast

# raw and tokenized files
SRC_RAW=${DATA_PATH}/train.$SRC
TGT_RAW=${DATA_PATH}/train.$TGT
SRC_TOK=$SRC_RAW.tok
TGT_TOK=$TGT_RAW.tok


# BPE / vocab files
BPE_CODES=$PROC_PATH/codes
SRC_VOCAB=$PROC_PATH/vocab.$SRC
TGT_VOCAB=$PROC_PATH/vocab.$TGT
FULL_VOCAB=$PROC_PATH/vocab.$SRC-$TGT
FULL_VOCAB_REV=$PROC_PATH/vocab.$TGT-$SRC

# train / valid / test monolingual BPE data
SRC_TRAIN_BPE=$PROC_PATH/train.$SRC
TGT_TRAIN_BPE=$PROC_PATH/train.$TGT
SRC_VALID_BPE=$PROC_PATH/valid.$SRC
TGT_VALID_BPE=$PROC_PATH/valid.$TGT
SRC_TEST_BPE=$PROC_PATH/test.$SRC
TGT_TEST_BPE=$PROC_PATH/test.$TGT

PARA_SRC_TRAIN_BPE=$PROC_PATH/train.$SRC-$TGT.$SRC
PARA_TGT_TRAIN_BPE=$PROC_PATH/train.$SRC-$TGT.$TGT
PARA_SRC_VALID_BPE=$PROC_PATH/valid.$SRC-$TGT.$SRC
PARA_TGT_VALID_BPE=$PROC_PATH/valid.$SRC-$TGT.$TGT
PARA_SRC_TEST_BPE=$PROC_PATH/test.$SRC-$TGT.$SRC
PARA_TGT_TEST_BPE=$PROC_PATH/test.$SRC-$TGT.$TGT

PARA_SRC_TRAIN_BPE_REV=$PROC_PATH/train.$TGT-$SRC.$SRC
PARA_TGT_TRAIN_BPE_REV=$PROC_PATH/train.$TGT-$SRC.$TGT
PARA_SRC_VALID_BPE_REV=$PROC_PATH/valid.$TGT-$SRC.$SRC
PARA_TGT_VALID_BPE_REV=$PROC_PATH/valid.$TGT-$SRC.$TGT
PARA_SRC_TEST_BPE_REV=$PROC_PATH/test.$TGT-$SRC.$SRC
PARA_TGT_TEST_BPE_REV=$PROC_PATH/test.$TGT-$SRC.$TGT

# valid / test file raw data
unset VALID_SRC_RAW VALID_TGT_RAW TEST_SRC_RAW TEST_TGT_RAW
VALID_SRC_RAW=${DATA_PATH}/valid.$SRC
VALID_TGT_RAW=${DATA_PATH}/valid.$TGT
TEST_SRC_RAW=${DATA_PATH}/test.$SRC
TEST_TGT_RAW=${DATA_PATH}/test.$TGT
VALID_SRC_TOK=${VALID_SRC_RAW}.tok
VALID_TGT_TOK=${VALID_TGT_RAW}.tok
TEST_SRC_TOK=${TEST_SRC_RAW}.tok
TEST_TGT_TOK=${TEST_TGT_RAW}.tok

#PARA_SRC_VALID=${DATA_PATH}/valid-ref.$SRC
#PARA_TGT_VALID=${DATA_PATH}/valid-ref.$TGT
#PARA_SRC_TEST=$PARA_PATH/test-ref.$SRC
#PARA_TGT_TEST=$PARA_PATH/test-ref.$TGT

echo "$SRC monolingual data concatenated in: $SRC_RAW"
echo "$TGT monolingual data concatenated in: $TGT_RAW"
ls $TGT_RAW
ls $SRC_RAW

# tokenize data
if ! [[ -f "$SRC_TOK" ]]; then
  echo "Tokenize $SRC monolingual data..."
  cat $SRC_RAW | $TOKENIZER -l $SRC -no-escape -threads $N_THREADS > $SRC_TOK
fi

if ! [[ -f "$TGT_TOK" ]]; then
  echo "Tokenize $TGT monolingual data..."
  cat $TGT_RAW | $TOKENIZER -l $TGT -no-escape -threads $N_THREADS > $TGT_TOK
fi
echo "$SRC monolingual data tokenized in: $SRC_TOK"
echo "$TGT monolingual data tokenized in: $TGT_TOK"

if [ ! -f "$BPE_CODES" ] && [ -f "$RELOAD_CODES" ]; then
  echo "Reloading BPE codes from $RELOAD_CODES ..."
  cp $RELOAD_CODES $BPE_CODES
fi

# learn BPE codes
if [ ! -f "$BPE_CODES" ]; then
  echo "Learning BPE codes..."
  $FASTBPE learnbpe $CODES $SRC_TOK $TGT_TOK > $BPE_CODES
fi
echo "BPE learned in $BPE_CODES"

# apply BPE codes
if ! [[ -f "$SRC_TRAIN_BPE" ]]; then
  echo "Applying $SRC BPE codes..."
  $FASTBPE applybpe $SRC_TRAIN_BPE $SRC_TOK $BPE_CODES
fi
if ! [[ -f "$TGT_TRAIN_BPE" ]]; then
  echo "Applying $TGT BPE codes..."
  $FASTBPE applybpe $TGT_TRAIN_BPE $TGT_TOK $BPE_CODES
fi
echo "BPE codes applied to $SRC in: $SRC_TRAIN_BPE"
echo "BPE codes applied to $TGT in: $TGT_TRAIN_BPE"

if [[ ! -f "$SRC_VOCAB" ]]; then
  $FASTBPE getvocab $SRC_TRAIN_BPE > $SRC_VOCAB
fi
echo "$SRC vocab in: $SRC_VOCAB"
if [[ ! -f "$TGT_VOCAB" ]]; then
  $FASTBPE getvocab $TGT_TRAIN_BPE > $TGT_VOCAB
fi
echo "$TGT vocab in: $TGT_VOCAB"

# reload full vocabulary
if [ ! -f "$FULL_VOCAB" ] && [ -f "$RELOAD_VOCAB" ]; then
  echo "Reloading vocabulary from $RELOAD_VOCAB ..."
  cp $RELOAD_VOCAB $FULL_VOCAB
fi

# extract full vocabulary
if ! [[ -f "$FULL_VOCAB" || -f "$FULL_VOCAB_REV" ]]; then
  echo "Extracting vocabulary..."
  $FASTBPE getvocab $SRC_TRAIN_BPE $TGT_TRAIN_BPE > $FULL_VOCAB
fi
echo "Full vocab in: $FULL_VOCAB"

cp $FULL_VOCAB $FULL_VOCAB_REV

# binarize data
if ! [[ -f "$SRC_TRAIN_BPE.pth" ]]; then
  echo "Binarizing $SRC data..."
  $MAIN_PATH/XLM/preprocess.py $FULL_VOCAB $SRC_TRAIN_BPE
fi
if ! [[ -f "$TGT_TRAIN_BPE.pth" ]]; then
  echo "Binarizing $TGT data..."
  $MAIN_PATH/XLM/preprocess.py $FULL_VOCAB $TGT_TRAIN_BPE
fi
echo "$SRC binarized data in: $SRC_TRAIN_BPE.pth"
echo "$TGT binarized data in: $TGT_TRAIN_BPE.pth"

echo "Tokenizing valid and test data..."
echo $DATA_PATH/test.$SRC

if ! [[ -f "$VALID_SRC_TOK" ]]; then
	cat $VALID_SRC_RAW | $TOKENIZER -l $SRC -no-escape -threads $N_THREADS > $VALID_SRC_TOK
fi
if ! [[ -f "$VALID_TGT_TOK" ]]; then
	cat $VALID_TGT_RAW | $TOKENIZER -l $TGT -no-escape -threads $N_THREADS > $VALID_TGT_TOK
fi
if ! [[ -f "$TEST_SRC_TOK" ]]; then
	cat $TEST_SRC_RAW | $TOKENIZER -l $SRC -no-escape -threads $N_THREADS > $TEST_SRC_TOK
fi
if ! [[ -f "$TEST_TGT_TOK" ]]; then
	cat $TEST_TGT_RAW | $TOKENIZER -l $TGT -no-escape -threads $N_THREADS > $TEST_TGT_TOK
fi
echo "Applying BPE to valid and test files..."
$FASTBPE applybpe $PARA_SRC_VALID_BPE $VALID_SRC_TOK $BPE_CODES $SRC_VOCAB
$FASTBPE applybpe $PARA_TGT_VALID_BPE $VALID_TGT_TOK $BPE_CODES $TGT_VOCAB
$FASTBPE applybpe $PARA_SRC_TEST_BPE  $TEST_SRC_TOK  $BPE_CODES $SRC_VOCAB
$FASTBPE applybpe $PARA_TGT_TEST_BPE  $TEST_TGT_TOK  $BPE_CODES $TGT_VOCAB

echo "Binarizing data..."
rm -f $PARA_SRC_VALID_BPE.pth $PARA_TGT_VALID_BPE.pth $PARA_SRC_TEST_BPE.pth $PARA_TGT_TEST_BPE.pth
$MAIN_PATH/XLM/preprocess.py $FULL_VOCAB $PARA_SRC_VALID_BPE
$MAIN_PATH/XLM/preprocess.py $FULL_VOCAB $PARA_TGT_VALID_BPE
$MAIN_PATH/XLM/preprocess.py $FULL_VOCAB $PARA_SRC_TEST_BPE
$MAIN_PATH/XLM/preprocess.py $FULL_VOCAB $PARA_TGT_TEST_BPE


#
# Link data
#
echo "$SRC_VALID_BPE.pth"
cp $PARA_SRC_VALID_BPE.pth $SRC_VALID_BPE.pth
cp $PARA_TGT_VALID_BPE.pth $TGT_VALID_BPE.pth
cp $PARA_SRC_TEST_BPE.pth  $SRC_TEST_BPE.pth
cp $PARA_TGT_TEST_BPE.pth  $TGT_TEST_BPE.pth
cp $SRC_TRAIN_BPE.pth $PARA_SRC_TRAIN_BPE.pth
cp $TGT_TRAIN_BPE.pth $PARA_TGT_TRAIN_BPE.pth

cp $PARA_SRC_VALID_BPE.pth $PARA_SRC_VALID_BPE_REV.pth
cp $PARA_TGT_VALID_BPE.pth $PARA_TGT_VALID_BPE_REV.pth
cp $PARA_SRC_TEST_BPE.pth  $PARA_SRC_TEST_BPE_REV.pth
cp $PARA_TGT_TEST_BPE.pth  $PARA_TGT_TEST_BPE_REV.pth
cp $PARA_SRC_TRAIN_BPE.pth $PARA_SRC_TRAIN_BPE_REV.pth
cp $PARA_TGT_TRAIN_BPE.pth $PARA_TGT_TRAIN_BPE_REV.pth

#
# Summary
#
echo ""
echo "===== Data summary"
echo "Monolingual training data:"
echo "    $SRC: $SRC_TRAIN_BPE.pth"
echo "    $TGT: $TGT_TRAIN_BPE.pth"
echo "Monolingual validation data:"
echo "    $SRC: $SRC_VALID_BPE.pth"
echo "    $TGT: $TGT_VALID_BPE.pth"
echo "Monolingual test data:"
echo "    $SRC: $SRC_TEST_BPE.pth"
echo "    $TGT: $TGT_TEST_BPE.pth"
echo "Parallel validation data:"
echo "    $SRC: $PARA_SRC_VALID_BPE.pth"
echo "    $TGT: $PARA_TGT_VALID_BPE.pth"
echo "Parallel test data:"
echo "    $SRC: $PARA_SRC_TEST_BPE.pth"
echo "    $TGT: $PARA_TGT_TEST_BPE.pth"
echo ""
