#!/usr/bin/env bash
#
# orthologer - Pipeline for deducing orthology from fasta files.
# Copyright (C) 2005-2023 E. Zdobnov lab: F. Tegenfeldt
# <fredrik.tegenfeldt@unige.ch> based on code by E. Kriventseva
#
# This file is part of orthologer.
#
# orthologer is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# orthologer is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with orthologer.  If not, see <http://www.gnu.org/licenses/>.
#
#--- check if environment is ok -----------------
[[ "${LOADED_COMMON_PIPELINE_SH+1}" == "" ]] && echo "ERROR $0 : main config has not been loaded" 1>&2 && exit 255
#------------------------------------------------

set -u

#----------------------------------------------
# setup alignment env variables
#----------------------------------------------
__ALIGNMENT_PREOPTS=${__ALIGNMENT_PREOPTS:=""}
__ALIGNMENT_ALIOPTS=${__ALIGNMENT_ALIOPTS:=""}
__ALIGNMENT_CNVOPTS=${__ALIGNMENT_CNVOPTS:=""}
__ALIGNMENT_TMPALI=${__ALIGNMENT_TMPALI:=""}
__ALIGNMENT_TMPQUERY=${__ALIGNMENT_TMPQUERY:=""}
setup_alignment_mmseqs() {
    # TARGET FORMAT:
    #    query   target  score evalue pid startQ endQ startD endD
    # set formatstr
    #       1      2      3      4      5     6        7      8       9       10   11     12   13   14     15
    local fmtstr="query,target,evalue,pident,raw,bits,alnlen,qcov,tcov,qlen,qstart,qend,tlen,tstart,tend"

    # set pre-filter opts
    ###__ALIGNMENT_PREOPTS="--max-seqs $ALIGNMENT_NUMALIGNMENTS -s $MMSEQS_SENSITIVITY"
    __ALIGNMENT_PREOPTS="--max-seq-len $ALIGNMENT_NUMALIGNMENTS -s $MMSEQS_SENSITIVITY"

    # main alignment opts
    __ALIGNMENT_ALIOPTS="$OPTS_MMSEQS -e $MAX_EVALUE"
    
    # convert opts
    __ALIGNMENT_CNVOPTS="--format-output $fmtstr"

    # create tmp files
    tmp_make __ALIGNMENT_TMPALI  mmseqsraw
    tmp_make __ALIGNMENT_TMPCNV  mmseqsali
    tmp_add_file "${__ALIGNMENT_TMPALI}.dbtype"
    tmp_add_file "${__ALIGNMENT_TMPALI}.index"
}

#----------------------------------------------
# run the alignment
#----------------------------------------------
run_alignment_mmseqs() {
    local ec=0
    #--------------------------------
    #--- 1. PROCESS ARGUMENTS     ---
    #--------------------------------
    [[ $# -eq 4 ]] || { error_msg "usage: $0 <query file> <DB file> <output> <nthreads>"; return 255; }
    local query="$1"
    local dbase="$2"
    local output="$3"
    local nthreads="$4"

    # make sure openmp does not use more that nthreads
    export OMP_NUM_THREADS=$nthreads

    # timing
    local startdate=$(date)
    local tA=$(time_msec)

    local querydb="${query}.${EXT_DBPRE}"
    local dbasedb="${dbase}.${EXT_DBPRE}"


    #--------------------------------
    #--- 2. PREFILTER            ---
    #--------------------------------
    info_msg "mmseqs: run prefilter"
    local mmpre=${output}.mmpre

    # run prefilter
    local tpref=$(time_msec)
    "$EXE_MMSEQS" prefilter "$querydb" "$dbasedb" "$mmpre" $__ALIGNMENT_PREOPTS --threads $nthreads > /dev/null #$tmplog
    if [[ $? -ne 0 ]]; then
	    #tmp_mark_keep $tmplog
        error_msg "alignment_mmseqs.sh: failed running prefilter on $(hostname)"
        error_msg "cmd: $EXE_MMSEQS prefilter $querydb $dbasedb $mmpre $__ALIGNMENT_PREOPTS --threads $nthreads"
        return 255
    fi
    local dTpref=$(delta_time $tpref)

    #--------------------------------
    #--- 3. RUN ALIGNMENT         ---
    #--------------------------------
    info_msg "mmseqs: run aligner"

    # run aligner and output to tmpfile
    local talign=$(time_msec)
    "$EXE_MMSEQS" align "$querydb" "$dbasedb" "$mmpre" "$__ALIGNMENT_TMPALI" $__ALIGNMENT_ALIOPTS --threads $nthreads > /dev/null # $tmplog
    ec=$?
    local dTalign=$(delta_time $talign)
    if [[ $ec -ne 0 ]]; then
        {
            error_msg "# Failed MMSEQS - alignment"
            echo "cmd: $EXE_MMSEQS align  $querydb $dbasedb $mmpre $__ALIGNMENT_TMPALI $__ALIGNMENT_ALIOPTS  --threads $nthreads"
            ls -l ${querydb}_h $querydb $dbasedb $mmpre $__ALIGNMENT_TMPALI
            echo "------------------"
            ###cat "$tmplog" 
            echo "------------------"
            error_msg "# Error     : $ec"
            error_msg "# Query     : $query"
            error_msg "# DB        : $dbase"
            error_msg "# Opts      : $__ALIGNMENT_ALIOPTS -d \"$dbase\" -i \"$query\""
            error_msg "# Output    : $SYS_UNAME:$__ALIGNMENT_TMPALI"
            error_msg "# Error log : ${output}.ali.err"
        } >& "${output}.ali.err"
        return 255
    fi

    # convert to blast like format
    info_msg "mmseqs: run convertalis"

    # set options with custom format
    local tconv=$(time_msec)
    "$EXE_MMSEQS" convertalis "$querydb" "$dbasedb" "$__ALIGNMENT_TMPALI" "$__ALIGNMENT_TMPCNV" $__ALIGNMENT_CNVOPTS --threads $nthreads > /dev/null
    if [[ $? -ne 0 ]]; then
        {
            error_msg "# Failed MMSEQS - convertalis"
            echo "cmd: $EXE_MMSEQS convertalis  $querydb $bdbasedb $__ALIGNMENT_TMPALI $__ALIGNMENT_TMPCNV $__ALIGNMENT_CNVOPTS --threads $nthreads"
            ls -l ${querydb}* ${dbasedb}* $mmpre $__ALIGNMENT_TMPALI $__ALIGNMENT_TMPCNV
            echo "------------------"
            ##        cat "$tmplog" 
            echo "------------------"
            error_msg "# Error     : $ec"
            error_msg "# Query     : $query"
            error_msg "# DB        : $dbase"
            error_msg "# Opts      : $opts"
            error_msg "# Input     : $SYS_UNAME:$__ALIGNMENT_TMPALI"
            error_msg "# Output    : $SYS_UNAME:$__ALIGNMENT_TMPCNV"
            error_msg "# Error log : ${output}.cnv.err"
        } >& "${output}.cnv.err"

        #    tmp_mark_keep "$__ALIGNMENT_TMPALI"
        #    tmp_mark_keep "$__ALIGNMENT_TMPCNV"
        #    tmp_mark_keep "$querydb"
        #    tmp_mark_keep "$dbasedb"
        return 255
    fi
    local dTconv=$(delta_time $tconv)
    local dT=$(delta_time $tA)
    #--------------------------------
    #--- 6. POST-PROCESSING       ---
    #--------------------------------
    # create correct format
    # BLAST LIKE OUTPUT (default)
    # Fields:  1      2               3       4       5            6       7       8       9       10      11              12
    #    query        target          pid/100 alnlen  missmatch    gapopen startQ  endQ    startD  endD    e-value         bit-score
    # TCAST:000015    AAEGL:006742    0.126   1577    1378         0       0       229     0       1577    8.29E-06        48
    #
    # CUSTOM OUTPUT
    #      1        2        3        4      5       6       7      8      9    10    11    12     13    14     15
    #    query   target   evalue   pident   raw   bits    alnlen   qcov   tcov   qlen qstart qend  tlen  tstart tend
    #
    # TARGET FORMAT:
    #    query   target  score evalue pid startQ endQ startD endD
    #
    info_msg "mmseqs: reformat output"

    ec=0
    #
    # Possibly rescale score (S->raw score, B->bit score)
    #  0 - S*alnlen/seqlen
    #  1 - <removed>
    #  2 - S/seqlen    (DEFAULT)
    #  3 - S
    #  4 - B
    #
    ####info_msg "mmseqs: score scale mode = $ALIGNMENT_SCORE_MODE"
    {
        print_header "mmseqs align $querydb $dbasedb $__ALIGNMENT_ALIOPTS" "$startdate"
        echo "# mmseqs align, without evalue rescaling"
        echo "# ALIGNMENT_SCORE_MODE = $ALIGNMENT_SCORE_MODE"
        echo "# MMSEQS_SENSITIVITY   = $MMSEQS_SENSITIVITY"
        common_fmt_header
        local trescore=$(time_msec)
        #if (( ALIGNMENT_SCORE_MODE == 0 )); then
        #    # S' = 2*S*alnlen/(qlen+tlen)
        #    grep -v "^#" "$__ALIGNMENT_TMPCNV"    | awk '{ print $1,$2,2.0*$5*$7/($10+$13),$3,$4,$11,$12,$14,$15; }'
        
        #elif (( ALIGNMENT_SCORE_MODE == 2 )); then
        # S' = 100*S/(query+target lengths) (the scaling is just to make the score a larger number)
        ### OLD grep -v "^#" "$__ALIGNMENT_TMPCNV"    | awk '{ print $1,$2,100.0*$5/($10+$13),$3,$4,$11,$12,$14,$15; }'

        # filter result and reformat
        # 1. remove lines starting with '#'
        # 2. remove alignments that are too short
        awk -v X="$MIN_OVERLAP" '!/^#/ && ($15-$14) >= X && ($12-$11) >= X { print $1,$2,100.0*$5/($10+$13),$3,$4,$11,$12,$14,$15; }' "$__ALIGNMENT_TMPCNV"
        
        #elif (( ALIGNMENT_SCORE_MODE == 3 )); then
        #    # S' = S
        #    grep -v "^#" "$__ALIGNMENT_TMPCNV"    | awk '{ print $1,$2,$5,$3,$4,$11,$12,$14,$15; }'
        
        #elif (( ALIGNMENT_SCORE_MODE == 4 )); then
        #    # use bit score
        #    grep -v "^#" "$__ALIGNMENT_TMPCNV"    | awk '{ print $1,$2,$6,$3,$4,$11,$12,$14,$15; }'
        #fi
        local dTrescore=$(delta_time $trescore)    
        echo "#TDETAIL: $dTpref $dTalign $dTconv $dTrescore $(hostname)"
        print_footer "$dT"
        
    } > "$output"

    #--------------------------------
    #--- 8. CLEANUP               ---
    #--------------------------------
    chmod 644 "$output"
}

#----------------------------------------------
# If the script has not been called before
# run the setup.
# Otherwise run the alignment, expecting arguments.
#----------------------------------------------
ALIGNMENT_EC=0
if [[ -z $__ALIGNMENT_PREOPTS ]]; then
    # setup call
    setup_alignment_mmseqs
else
    # run alignment - check args
    run_alignment_mmseqs "$@"
    ALIGNMENT_EC=$?
fi
