# Author: Daniel Ortiz Mart\'inez
# *- bash -*

# Trains a phrase-based model given a parallel corpus.

# \textbf{Categ}: modelling

########
print_desc()
{
    echo "thot_tm_train written by Daniel Ortiz"
    echo "thot_tm_train trains a translation model from a bilingual corpus"
    echo "type \"thot_tm_train --help\" to get usage information"
}

########
version()
{
    echo "thot_tm_train is part of the thot package"
    echo "thot version "${version}
    echo "thot is GNU software written by Daniel Ortiz"
}

########
usage()
{
    echo "thot_tm_train           [-pr <int>]"
    echo "                        -s <string> -t <string> {-o <string>|-a <string>}"
    echo "                        [-g <string>] [-dict] [-nit <int>] [-af <float>]"
    echo "                        [-cpr <float>] [-np <float>] [-m <int>] [-ao <string>]"
    echo "                        [-to <int>] [-unk]"
    if [ ! -z "${LDB_CXX}" ]; then
        echo "                        [-bdb]"
    fi
    if [ ! -z "${LEVELDB_LIB}" ]; then
        echo "                        [-ldb]"
    fi
    echo "                        [-qs <string>] [-tdir <string>]"
    echo "                        [-sdir <string>] [-debug] [--help] [--version]"
    echo ""
    echo "-pr <int>               Number of processors (1 by default)"
    echo "-s <string>             File with source sentences"
    echo "-t <string>             File with target sentences"
    echo "-o <string>             Output directory common to all processors"
    echo "-a <string>             Output directory previously generated by thot_tm_train."
    echo "                        The new model will be added to models already generated"
    echo "-g <string>             Use external alignment file in GIZA format for phrase"
    echo "                        extraction. Thot's native alignment file is not used" 
    echo "-dict                   Input data is considered as a dictionary, so the"
    echo "                        sentence pairs are introduced as phrase table entries."
    echo "                        This option cannot be combined with -g."
    echo "-nit <int>              Number of iterations of the EM algorithm executed by"
    echo "                        the thot_gen_sw_model tool (5 by default)"
    echo "-af <float>             Alignment smoothing interpolation factor for"
    echo "                        single-word models"
    echo "-cpr <float>            Pruning parameter used during the estimation of single"
    echo "                        word alignment models (0.00001 by default)"
    echo "-np <float>             Probability assigned to the alignment with the NULL"
    echo "                        word for single-word models"
    echo "-m <int>                Maximum target phrase length during phrase model"
    echo "                        estimation (10 by default)"
    echo "-ao <string>            Operation between alignments to be executed"
    echo "                        (and|or|sum|sym1|sym2|grd)."
    echo "-to <int>               Maximum number of translation options for each target"
    echo "                        phrase that are considered during a translation process"
    echo "                        (20 by default)"
    echo "-unk                    Introduce special unknown word symbol during"
    echo "                        estimation"
    if [ ! -z "${LDB_CXX}" ]; then
        echo "-bdb                    Generate on-disk phrase table in BDB format"
    fi
    if [ ! -z "${LEVELDB_LIB}" ]; then
        echo "-ldb                    Generate on-disk phrase table in LevelDB format"
    fi
    echo "-qs <string>            Specific options to be given to the qsub"
    echo "                        command (example: -qs \"-l pmem=1gb\")"
    echo "                        NOTES:"
    echo "                         a) ignore this if not using a PBS cluster"
    echo "                         b) -qs option may be crucial to ensure the correct"
    echo "                            execution of the tool. The main purpose of -qs"
    echo "                            is to reserve the required cluster resources."
    echo "                            If the necessary resources are not met the"
    echo "                            execution will abort."
    echo "                            Resources are reserved in different ways depending"
    echo "                            on the cluster software. For instance, if using SGE"
    echo "                            software, -qs \"-l h_vmem=1G,h_rt=10:00:00\","
    echo "                            requests 1GB of virtual memory and a time limit"
    echo "                            of 10 hours" 
    echo "-tdir <string>          Directory for temporary files (/tmp by default)."
    echo "                        NOTES:"
    echo "                         a) give absolute paths when using pbs clusters"
    echo "                         b) ensure there is enough disk space in the partition"
    echo "-sdir <string>          Absolute path of a directory common to all"
    echo "                        processors. If not given, \$HOME will be used."
    echo "                        NOTES:"
    echo "                         a) give absolute paths when using pbs clusters"
    echo "                         b) ensure there is enough disk space in the partition"
    echo "-debug                  After ending, do not delete temporary files"
    echo "                        (for debugging purposes)"
    echo "--help                  Display this help and exit"
    echo "--version               Output version information and exit"
}

########
is_absolute_path()
{
    case $1 in
        /*) echo 1 ;;
        *) echo 0 ;;
    esac
}

########
get_absolute_path()
{
    file=$1
    # Check if an absolute path was given
    absolute=`is_absolute_path $file`
    if [ $absolute -eq 1 ]; then
        echo $file
    else
        oldpwd=$PWD
        basetmp=`"$BASENAME" "$PWD/$file"`
        dirtmp=`"$DIRNAME" "$PWD/$file"`
        cd "$dirtmp"
        result=${PWD}/${basetmp}
        cd "$oldpwd"
        echo $result
    fi
}

########
create_desc_file()
{
    # Determine model type
    if [ ! -z "${LDB_CXX}" -a ${bdb_given} -eq 1 ]; then
        modeltype="\$(${LIBDIR_VARNAME})/fast_bdb_phrase_model_factory.so"
    else
        if [ ! -z "${LEVELDB_LIB}" -a ${ldb_given} -eq 1 ]; then
            modeltype="\$(${LIBDIR_VARNAME})/leveldb_phrase_model_factory.so"
        else
            modeltype="\$(${LIBDIR_VARNAME})/incr_phrase_model_factory.so"
        fi
    fi

    # Create descriptor file
    if [ ${o_given} -eq 1 ]; then
        # -o option was given
        echo "thot tm descriptor # tool: thot_tm_train" > "${outd}/tm_desc"
        echo "${modeltype} ${relative_prefix} main # source file: ${scorpus} ; target file: ${tcorpus}" >> "${outd}/tm_desc"
    else
        # -a option was given
        echo "${modeltype} ${relative_prefix} ${outsubdir} # source file: ${scorpus} ; target file: ${tcorpus}" >> "${outd}/tm_desc"
    fi
}

########
generate_outsubdir_name()
{
    if [ ${o_given} -eq 1 ]; then
        # -o option was given
        echo "main"
    else
        # -a option was given
        success=0
        for num in 1 2 3 4 5 6 7 8 9 10; do
            if [ ! -d "${outd}/additional_${num}" ]; then
                echo "additional_${num}"
                success=1
                break
            fi
        done

        if [ ${success} -eq 0 ]; then
            # Maximum number of models was exceeded
            echo "Error! a maximum of 10 additional models are allowed" >&2
            return 1
        fi
    fi
}

########
if [ $# -lt 1 ]; then
    print_desc
    exit 1
fi

# Read parameters
pr_given=0
pr_val=1
s_given=0
t_given=0
o_given=0
a_given=0
g_given=0
nit_given=0
niters=5
af_given=0
cpr_given=0
cpr_val=0.00001
np_given=0
m_val=10
ao_given=0
ao_opt="-ao sym1"
#ao_opt="-ao grd"
to_given=0
to_val=20
dict_given=0
unk_given=0
qs_given=0
bdb_given=0
ldb_given=0
tdir_given=0
tdir="/tmp"
sdir_given=0
sdir=$HOME
debug=0

while [ $# -ne 0 ]; do
    case $1 in
        "--help") usage
            exit 0
            ;;
        "--version") version
            exit 0
            ;;
        "-pr") shift
            if [ $# -ne 0 ]; then
                pr_val=$1
                pr_given=1
            fi
            ;;
        "-s") shift
            if [ $# -ne 0 ]; then
                scorpus=$1
                s_given=1
            fi
            ;;
        "-t") shift
            if [ $# -ne 0 ]; then
                tcorpus=$1
                t_given=1
            fi
            ;;
        "-o") shift
            if [ $# -ne 0 ]; then
                outd=$1
                o_given=1
            fi
            ;;
        "-a") shift
            if [ $# -ne 0 ]; then
                outd=$1
                a_given=1
            fi
            ;;
        "-g") shift
              if [ $# -ne 0 ]; then
                  gfile=$1
                  g_given=1
                  g_opt="-g ${gfile}"
              fi
              ;;
        "-to") shift
            if [ $# -ne 0 ]; then
                to_val=$1
                to_given=1
            fi
            ;;
        "-nit") shift
            if [ $# -ne 0 ]; then
                niters=$1
                nit_given=1
            fi
            ;;
        "-af") shift
            if [ $# -ne 0 ]; then
                af_opt="-af $1"
                af_given=1
            else
                af_given=0
            fi
            ;;
        "-cpr") shift
            if [ $# -ne 0 ]; then
                 cpr_opt="-cpr $1"
                 cpr_given=1
            fi
            ;;
        "-np") shift
            if [ $# -ne 0 ]; then
                np_opt="-np $1"
                np_given=1
            else
                np_given=0
            fi
            ;;
        "-m") shift
            if [ $# -ne 0 ]; then
                m_val=$1
                m_given=1
            fi
            ;;
        "-ao") shift
            if [ $# -ne 0 ]; then
                ao_opt="-ao $1"
                ao_given=1
            fi
            ;;
        "-qs") shift
            if [ $# -ne 0 ]; then
                qs_opt="-qs"
                qs_par="$1"
                qs_given=1
            else
                qs_given=0
            fi
            ;;
        "-dict") dict_given=1
            dict_opt="-dict"
            ;;
        "-unk") unk_given=1
            unk_opt="-unk"
            ;;
        "-bdb") bdb_given=1
            ;;
        "-ldb") ldb_given=1
            ;;
        "-tdir") shift
            if [ $# -ne 0 ]; then
                tdir=$1
                tdir_given=1
            fi
            ;;
        "-sdir") shift
            if [ $# -ne 0 ]; then
                sdir=$1
                sdir_given=1
            fi
            ;;
        "-debug") debug=1
            debug_opt="-debug"
            ;;
    esac
    shift
done

# Check parameters
if [ ${s_given} -eq 0 ]; then
    echo "Error! -s parameter not given" >&2
    exit 1
else
    if [ ! -f "${scorpus}" ]; then
        echo "Error! file ${scorpus} does not exist" >&2
        exit 1
    else
        # Obtain absolute path
        scorpus=`get_absolute_path "$scorpus"`
    fi
fi

if [ ${t_given} -eq 0 ]; then        
    echo "Error! -t parameter not given" >&2
    exit 1
else
    if [ ! -f "${tcorpus}" ]; then
        echo "Error! file ${tcorpus} does not exist" >&2
        exit 1
    else
        # Obtain absolute path
        tcorpus=`get_absolute_path "$tcorpus"`
    fi
fi



# Check that source and target files are parallel
nl_source=`wc -l "$scorpus" | "$AWK" '{printf"%d",$1}'`
nl_target=`wc -l "$tcorpus" | "$AWK" '{printf"%d",$1}'`

if [ "${nl_source}" -ne "${nl_target}" ]; then
    echo "Error! source and target files have not the same number of lines" >&2 
    exit 1
fi

if [ ${o_given} -eq 0 -a ${a_given} -eq 0 ]; then
    echo "Error! -o or -a parameter not given!" >&2    
    exit 1
fi

if [ ${o_given} -eq 1 -a ${a_given} -eq 1 ]; then
    echo "Error! -o or -a parameter cannot be given simultaneously!" >&2    
    exit 1
fi

if [ ${o_given} -eq 1 ]; then
    if [ -d "${outd}/main" ]; then
        echo "Warning! output directory does exist" >&2
    fi
    # Obtain absolute path
    outd=`get_absolute_path "$outd"`
fi

if [ ${a_given} -eq 1 ]; then
    if [ ! -d "${outd}/main" ]; then
        echo "Error! previous model estimated with thot_tm_train does not exist" >&2 
        exit 1
    fi
    # Obtain absolute path
    outd=`get_absolute_path "$outd"`
fi

if [ ${g_given} -eq 1 ]; then
    if [ ! -f "${gfile}" ]; then        
        echo "Error! file ${gfile} does not exist" >&2
        exit 1
    fi
fi

if [ ${tdir_given} -eq 1 ]; then
    if [ ! -d "${tdir}" ]; then
        echo "Error! directory ${tdir} does not exist" >&2
        exit 1   
    fi         
fi

if [ ${sdir_given} -eq 1 ]; then
    if [ ! -d "${sdir}" ]; then
        echo "Error! directory ${sdir} does not exist" >&2
        exit 1            
    fi
fi

if [ ${bdb_given} -eq 1 -a ${ldb_given} -eq 1 ]; then
    echo "Error! -bdb and -ldb options cannot be given simultaneously" >&2
    exit 1            
fi

# Obtain name of the output subdirectory where the language model will
# be stored
outsubdir=`generate_outsubdir_name` || exit 1

# Create output subdirectory
mkdir -p "${outd}/${outsubdir}" || { echo "Error! cannot create output directory" >&2; exit 1; }

# Train model
prefix="${outd}/${outsubdir}/src_trg"
relative_prefix="${outsubdir}/src_trg"
"${bindir}"/thot_pbs_gen_batch_phr_model -pr ${pr_val} \
    -s "$tcorpus" -t "$scorpus" -o $prefix "${g_opt}" -nit $niters "${af_opt}" "${cpr_opt}" \
    "${np_opt}" -m ${m_val} "${ao_opt}" -to ${to_val} "${dict_opt}" "${unk_opt}" \
    "${qs_opt}" "${qs_par}" -T "$tdir" -sdir "$sdir" ${debug_opt} || exit 1

# Process -bdb option if given
if [ ! -z "${LDB_CXX}" -a ${bdb_given} -eq 1 ]; then
    echo "* Generating on-disk phrase table in BDB format..." >&2
    echo "" >&2
    "${bindir}"/thot_gen_fbdb_ttable -p "$prefix" -o "$prefix" 2> "${prefix}.fbdb_err" || exit 1 
fi

# Process -ldb option if given
if [ ! -z "${LEVELDB_LIB}" -a ${ldb_given} -eq 1 ]; then
    echo "* Generating on-disk phrase table in LevelDB format..." >&2
    echo "" >&2
    "${bindir}"/thot_gen_leveldb_ttable -p "$prefix" -o "$prefix" 2> "${prefix}.ldb_err" || exit 1 
fi

# Create descriptor file
echo "* Generating descriptor file... " >&2
create_desc_file "$outd"
echo "" >&2
