#!/bin/bash

if [ $# -ne 4 ]; then
    echo "usage: $0 <input_dir/> <langs-string> <output_dir/> <step_name>"
    echo "  input_dir   : data/how2_300h/ or data/how2_153h/ ..."
    echo "  langs-string: en-pt or en or pt"
    echo "  output_dir  : path to out dir where processed data will be saved"
    echo "  step_name   : tokenize, joint-bpe"
    exit
fi

input_dir=`realpath $1`
lang_string=$2

mkdir -p $3
output_dir=`realpath $3`

nbpe=8000

echo "- nbpe: ${nbpe}"
sleep 3

step=$4

set_names=(train valid test)

IFS='-' read -r -a languages <<< ${lang_string}

FASTBPE=tools/fastBPE/fast

set -e

# General steps for data prep

# 1. Tokenize using tools/tokenize.sh
if [ "${step}" == "tokenize" ]; then

    for lang in ${languages[@]}; do
        for set_name in ${set_names[@]}; do

            in_file=${input_dir}/${lang}.${set_name}
            if [ -f ${in_file} ]; then
                echo "Tokenizing .. "${in_file}
                cat ${in_file} | ./tools/tokenize.sh ${lang} > ${in_file}.tok
                echo " .. tokenized file saved as "${in_file}.tok
            else
                echo "File not found: "${in_file}
                exit;
            fi

        done
    done

# 2. for joint MLM/TLM training, concat mono text from all the languages
elif [ "${step}" == "joint-bpe" ]; then

    mkdir -p ${output_dir}/preprocessed/
    combined_file=${output_dir}/preprocessed/xlm_${lang_string}.bpe.train
    if [ -f ${combined_file} ]; then
        echo "out file exists `wc -l ${combined_file}`"
        rm -i ${combined_file}
    fi

    set_name=${set_names[0]}
    for lang in ${languages[@]}; do
        in_file=${input_dir}/${lang}.${set_name}.tok
        if [ -f ${in_file} ]; then
            shuf ${in_file} >> ${combined_file}
        else
            echo "File not found: "${in_file}
            exit;
        fi
    done

    wc -l ${combined_file}

    # learn BPE
    codes_file=${output_dir}/preprocessed/codes
    ${FASTBPE} learnbpe ${nbpe} ${combined_file} > ${codes_file}

    # apply BPE
    for lang in ${languages[@]}; do
        for set_name in ${set_names[@]}; do

            in_file=${input_dir}/${lang}.${set_name}.tok

            ${FASTBPE} applybpe ${output_dir}/preprocessed/${set_name}.${lang} ${in_file} ${codes_file}

        done
    done

    # get post-BPE vocab
    ${FASTBPE} applybpe ${combined_file}.bpe ${combined_file} ${codes_file}

    vocab_file=${output_dir}/preprocessed/vocab
    cat ${combined_file}.bpe | ${FASTBPE} getvocab - > ${vocab_file}
    wc -l ${vocab_file}

    # binarize
    for lang in ${languages[@]}; do
        for set_name in ${set_names[@]}; do

            python preprocess.py ${vocab_file} ${output_dir}/preprocessed/${set_name}.${lang}

        done
    done


# 3. for mono MLM training, concat mono text from single language
elif [ "${step}" == "mono-bpe" ]; then

    mkdir -p ${output_dir}/preprocessed/
    mono_file=${output_dir}/preprocessed/mlm_${lang_string}.bpe.train
    if [ -f ${mono_file} ]; then
        echo "(mono-bpe) Out file exists `wc -l ${mono_file}`"
        rm -i ${mono_file}
    fi

    lang=${lang_string}
    set_name=${set_names[0]}
    in_file=${input_dir}/${lang}.${set_name}.tok
    if [ -f ${in_file} ]; then
        shuf ${in_file} >> ${mono_file}
    else
        echo "(mono-bpe) File not found: "${in_file}
        exit;
    fi


    wc -l ${mono_file}

    # learn BPE
    codes_file=${output_dir}/preprocessed/mono_${lang_string}_codes
    ${FASTBPE} learnbpe ${nbpe} ${mono_file} > ${codes_file}

    # apply BPE
    for set_name in ${set_names[@]}; do

        in_file=${input_dir}/${lang}.${set_name}.tok

        ${FASTBPE} applybpe ${output_dir}/preprocessed/${set_name}.${lang} ${in_file} ${codes_file}

    done

    # get post-BPE vocab
    ${FASTBPE} applybpe ${mono_file}.bpe ${mono_file} ${codes_file}

    vocab_file=${output_dir}/preprocessed/vocab
    cat ${mono_file}.bpe | ${FASTBPE} getvocab - > ${vocab_file}
    wc -l ${vocab_file}

    # binarize
    for lang in ${languages[@]}; do
        for set_name in ${set_names[@]}; do

            python preprocess.py ${vocab_file} ${output_dir}/preprocessed/${set_name}.${lang}

        done
    done

elif [ "${step}" == "soft-link" ]; then

    for lang in ${languages[@]}; do
        for s in ${set_names[@]}; do
            ln -sv ${output_dir}/preprocessed/${s}.${lang}.pth ${output_dir}/preprocessed/${s}.en-pt.${lang}.pth
        done
    done

else

    echo ${step}" not understood."
fi
