#!/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

#--------------------------------
#--- 1. PROCESS ARGUMENTS     ---
#--------------------------------
# get args
[[ $# -eq 4 ]] || fatal_msg "usage: $0 <query file> <DB file> <output> <nthreads>"
query="$1"
dbase="$2"
output="$3"
nthreads="$4"
(( nthreads > 0 )) || nthreads=1

###nthreads=3

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

# disable evalrescale as the score is not the true score
do_evalrescale=0

#--------------------------------
#--- 2. CHECK INPUT           ---
#--------------------------------
#HERE
#check_file     "$0" "$query"       || exit 255
#check_file     "$0" "$dbase"       || exit 255

#--------------------------------
#--- 3. CHECK REQUIREMENTS    ---
#--------------------------------
# check that formatdb files are there
#HERE
###check_formatdb_files "$query" 0 || exit 255
###check_formatdb_files "$dbase" 0 || exit 255

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

# get effective DB size
effdbsz=0
if (( do_evalrescale==1 )); then
    effdbsz=$(get_effdbsz)
    if (( effdbsz > 0 )); then
        cmd="$DIR_BRHCLUS/evalrescale"
        [[ -x $cmd ]] || fatal_msg "alignment_mmseqs: missing evalrescale tool $cmd"
    fi
else
    info_msg "mmseqs: NOTE e-value rescaling is disabled (the score is renormalised)"
fi

# get dir of query - needed further down...
###dirdb=$(dirname "$querydb")

# log file used in all steps
tmplog=""
tmp_make tmplog  mmseqslog || exit 255


#--------------------------------
#--- 3. SETUP SEARCH CMD   ---
#--------------------------------
maxeval=$MAX_EVALUE
dbntot=0
# if effective DB size, then correct MAXEVAL such that the cutoff
# corresponds to final e-value
if (( effdbsz > 0 )); then
    fdbStats="${dbase}.stats"
    [[ -f "$fdbStats" ]] || fatal_msg "alignment_mmseqs.sh: unable to find DB stats file $fdbStats"
    dbntot=$(grep '#NTOT:' "$fdbStats" 2>/dev/null | cut -d ':' -f 2)
    fdb="1.0"
    (( dbntot > 0 )) && fdb=$(calc "$effdbsz/$dbntot")
    maxeval=$(calc "$maxeval/$fdb")
fi

# set opts
opts="$OPTS_MMSEQS --threads $nthreads -e $maxeval"
# --rescore-mode 3 -> rescore global ali
# -s 5.7  -> lower value -> lower sensitivity , in range [1,7.5]
opts="$opts --max-accept $ALIGNMENT_NUMALIGNMENTS --rescore-mode $MMSEQS_RESCORE_MODE -s $MMSEQS_SENSITIVITY --start-sens $MMSEQS_SENSITIVITY"



#--------------------------------
#--- 4. RUN ALIGNMENT         ---
#--------------------------------
#info_msg "---------------------------------"
info_msg "mmseqs: run search"
#info_msg "---------------------------------"

# run aligner and output to tmpfile
tmp_aliraw=""
tmp_alicnv=""
tmp_make tmp_aliraw  mmseqs
tmp_make tmp_alicnv  mmseqsali

tmp_root=
tmp_base tmp_root
tmp_dir=$(mktemp -d $tmp_root/mmseqs_search_XXXXXX)

cmd="$EXE_MMSEQS search $querydb $dbasedb $tmp_aliraw $tmp_dir $opts --alignment-mode 3"
"$EXE_MMSEQS" search "$querydb" "$dbasedb" "$tmp_aliraw" "$tmp_dir" $opts --alignment-mode 3 >& $tmplog
ec=$?
rm -Rf $tmp_dir
if [[ $ec -ne 0 ]]; then
    {
        error_msg "# Failed MMSEQS - alignment"
        echo "cmd: $cmd"
        ls -l ${querydb}_h $querydb $dbasedb $tmp_aliraw
        echo "------------------"
        cat "$tmplog" 
        echo "------------------"
        error_msg "# Error     : $ec"
        error_msg "# Query     : $query"
        error_msg "# DB        : $dbase"
        error_msg "# Opts      : $opts -d \"$dbase\" -i \"$query\""
        error_msg "# Output    : $SYS_UNAME:$tmp_aliraw"
        error_msg "# Error log : ${output}.ali.err"
    } >& "${output}.ali.err"
    if [[ $nthreads -lt 2 ]]; then
        tmp_mark_keep "$tmp_aliraw"
    fi
    tmp_mark_keep "$query"
    tmp_mark_keep "$dbase"
    exit 255
fi

if [[ $nthreads -gt 1 ]]; then
    for (( i=0; i<$nthreads; i++ )); do
        tmp_add_file "${tmp_aliraw}.$i"
    done
fi

tmp_add_file "${tmp_aliraw}.dbtype"
tmp_add_file "${tmp_aliraw}.index"

[[ -f $tmp_aliraw ]] || fatal_msg "lost the temp file during run - check automatic cleanup of temp file!"


# custom format
opts="--format-output query,target,evalue,pident,raw,bits,alnlen,qcov,tcov,qlen,qstart,qend,tlen,tstart,tend"


# SKIP THIS CHECK FOR OLD MMSEQS
# convert to blast tab format
# TODO: convert to custom format
#use_blast_fmt="0"
#"$EXE_MMSEQS" convertalis help 2>&1 | grep -q 'format-output'
#if [[ $? -ne 0 ]]; then
#    use_blast_fmt="1"
#fi
#
#if [[ $use_blast_fmt == "0" ]]; then
#    info_msg "using new custom output format"
#    opts="--format-output query,target,evalue,pident,raw,alnlen,qcov,tcov,qlen,qstart,qend,tlen,tstart,tend"
#else
#    info_msg "using old BLAST like format"
#fi



# 
#cd $dirdb
#bquerydb=$(basename $querydb)
#bdbasedb=$(basename $dbasedb)
#info_msg "---------------------------------"
info_msg "mmseqs: run convertalis"
#info_msg "---------------------------------"

opts="$opts --threads $nthreads"
"$EXE_MMSEQS" convertalis "$querydb" "$dbasedb" "$tmp_aliraw" "$tmp_alicnv" $opts >& "$tmplog"
if [[ $? -ne 0 ]]; then
    {
        error_msg "# Failed MMSEQS - convertalis"
        echo "cmd: $EXE_MMSEQS convertalis  $querydb $bdbasedb $tmp_aliraw $tmp_alicnv $opts"
        ls -l ${querydb}* ${dbasedb}* $tmp_aliraw $tmp_alicnv
        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:$tmp_aliraw"
        error_msg "# Output    : $SYS_UNAME:$tmp_alicnv"
        error_msg "# Error log : ${output}.cnv.err"
    } >& "${output}.cnv.err"

    tmp_mark_keep "$tmp_aliraw"
    tmp_mark_keep "$tmp_alicnv"
    tmp_mark_keep "$querydb"
    tmp_mark_keep "$dbasedb"
    exit 255
fi

#=> TEMP CODE <==
#tmp_mark_keep "$tmp_aliraw"
#tmp_mark_keep "$tmp_alicnv"

#--------------------------------
#--- 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 "---------------------------------"
info_msg "mmseqs: reformat output"
#info_msg "---------------------------------"

tmp_make tmpali reformat
ec=0

echo "mmseqs score scale mode: $MMSEQS_SCORE_MODE"

if (( ALIGNMENT_SCORE_MODE == 0 )); then
    # S' = 2*S*alnlen/(qlen+tlen)
    grep -v "^#" "$tmp_alicnv"    | awk '{ print $1,$2,2.0*$5*$7/($10+$13),$3,$4,$11,$12,$14,$15; }' > "$tmpali"

elif (( ALIGNMENT_SCORE_MODE == 2 )); then
    # S' = 100*S/(query+target lengths) (the scaling is just to make the score a larger number)
    grep -v "^#" "$tmp_alicnv"    | awk '{ print $1,$2,100.0*$5/($10+$13),$3,$4,$11,$12,$14,$15; }' > "$tmpali"
    
elif (( ALIGNMENT_SCORE_MODE == 3 )); then
    # S' = S
    grep -v "^#" "$tmp_alicnv"    | awk '{ print $1,$2,$5,$3,$4,$11,$12,$14,$15; }' > "$tmpali"
    
elif (( ALIGNMENT_SCORE_MODE == 4 )); then
    # use bit score
    grep -v "^#" "$tmp_alicnv"    | awk '{ print $1,$2,$6,$3,$4,$11,$12,$14,$15; }' > "$tmpali"
fi

[[ $ec -eq 0 ]] || { tmp_mark_keep "$tmpali"; tmp_mark_keep "$tmp_alicnv"; fatal_msg "odd - failed reformatting alignment file on $SYS_UNAME, file $tmp_alicnv $tmpali"; }

# only rescale if effdbsz > 0
tmpres="$tmpali"
if (( effdbsz > 0 )); then
    
#    info_msg "---------------------------------"
    info_msg "mmseqs: rescale e-values"
#    info_msg "---------------------------------"
    fqueryLen="${query}.seqlen"

    cmdresc="$DIR_BRHCLUS/evalrescale"
    [[ -x $cmdresc ]] || fatal_msg "alignment_mmseqs: missing evalrescale tool $cmdresc"
    info_msg "cmd: $cmdresc  -f $tmpali -q $fqueryLen -s $effdbsz -n $dbntot -e $maxeval"
    
    tmp_make tmpres rescale
    "$cmdresc" -f "$tmpali" -q "$fqueryLen" -s $effdbsz -n $dbntot -e $maxeval >& $tmpres
    if [[ $? -ne 0 ]]; then
        tmp_mark_keep "$tmp_aliraw"
        tmp_mark_keep "$tmp_alicnv"
        tmp_mark_keep "$query"
        tmp_mark_keep "$dbase"
        tmp_mark_keep "$tmpali"
        tmp_mark_keep "$tmpres"
        error_msg "query     : $query"
        error_msg "dbase     : $dbase"
        error_msg "tmp_aliraw: $tmp_aliraw"
        error_msg "tmp_alicnv  : $tmp_alicnv"
        error_msg "tmpali    : $tmpali"
        error_msg "tmpres    : $tmpres"
        fatal_msg "odd - failed rescaling e-values in alignment file on $SYS_UNAME"
    fi
else
    info_msg "mmseqs: no evalue rescaling"
fi
dT=$(delta_time $tA)

#--------------------------------
#--- 7. ADD HEADER/FOOTER     ---
#--------------------------------
add_to_top "$tmpres" "$tmpres" "# mmseqs search MMSEQS_SCORE_MODE=$MMSEQS_SCORE_MODE\n#FMT id1 id2 score evalue pid start1 end1 start2 end2"          || fatal_msg "failed adding mmseqs header to top of file!"
#add_to_top "$tmpres" "$tmpres" "$(common_fmt_header)"    || fatal_msg "failed adding header to top of file!"

add_header_footer $0 "$tmpres" "$output" "$startdate" "$dT" "$cmd" || {
    tmp_mark_keep "$tmpres"
    fatal_msg "failed adding final header and footer of file!"
}

# copy file to output
#cp -f "$tmpres" "$output" || { tmp_mark_keep "$tmpres"; fatal_msg "failed copying mmseqs output $tmpres to $output"; }

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