#!/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_OPTS=${__ALIGNMENT_OPTS:=""}
__ALIGNMENT_TMPALI=${__ALIGNMENT_TMPALI:=""}
__ALIGNMENT_TMPQUERY=${__ALIGNMENT_TMPQUERY:=""}
setup_alignment_diamond() {
    # 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="qseqid sseqid evalue pident score bitscore length qcovhsp scovhsp qlen qstart qend slen sstart send"

    # if multithreading is used, this is where temp files are stored by diamond
    local parallel_tmp=
    tmp_base parallel_tmp

    # set opts
    __ALIGNMENT_OPTS="$OPTS_DIAMOND -e $MAX_EVALUE --masking none --parallel-tmpdir $parallel_tmp"

    # restrict nr of alignments - may also use --top (percentage of max score)
    __ALIGNMENT_OPTS="$__ALIGNMENT_OPTS --max-target-seqs $ALIGNMENT_NUMALIGNMENTS"

    # set output format - will be tab separated
    __ALIGNMENT_OPTS="$__ALIGNMENT_OPTS --outfmt 6 $fmtstr"

    # create tmp files
    tmp_make __ALIGNMENT_TMPALI    diamond_tmpalign
    tmp_make __ALIGNMENT_TMPQUERY  diamond_tmpquery
}

#----------------------------------------------
# run the alignment
#----------------------------------------------
run_alignment_diamond() {
    [[ $# -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


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

    #--------------------------------
    #--- 1. RUN ALIGNMENT         ---
    #--------------------------------
    info_msg "diamond: run aligner"

    # run aligner and output to tmpfile
    local ec=0

    # remove comments
    grep -v '^#' "$query" > "$__ALIGNMENT_TMPQUERY"
    info_msg "diamond command:  $EXE_DIAMOND blastp -d ${dbase}.dmnd -q $query -o $__ALIGNMENT_TMPALI $__ALIGNMENT_OPTS --threads $nthreads"
    ####<(grep -v '^#' "$query")
    local Talign=$(time_msec)
    $EXE_DIAMOND blastp -d "${dbase}.dmnd" -q "$__ALIGNMENT_TMPQUERY" -o "$__ALIGNMENT_TMPALI" --threads $nthreads $__ALIGNMENT_OPTS --quiet > /dev/null
    ec=$?
    local dTalign=$(delta_time $Talign)

    if [[ $ec -ne 0 ]]; then
        cp ${output}.ali.err ./
        error_msg "diamond failed - see $(pwd)/$(basename ${output}.ali.err)"
        {
            error_msg "# Failed DIAMOND - alignment"
            error_msg "# Error     : $ec"
            error_msg "# Query     : $query"
            error_msg "# DB        : $dbase"
            error_msg "# Opts      : $__ALIGNMENT_OPTS"
            error_msg "# Output    : $SYS_UNAME:$__ALIGNMENT_TMPALI"
            error_msg "# Error log : $(pwd)/${output}.ali.err"
        }
        #>& "${output}.ali.err"
        return 255
    fi

    [[ -f $__ALIGNMENT_TMPALI ]] || {
        error_msg "lost the temp file during run - check automatic cleanup of temp file!"
        return 255
    }

    if [[ ! -s $__ALIGNMENT_TMPALI ]]; then
        info_msg "diamond output is empty!"
        echo "#EMPTY" > $__ALIGNMENT_TMPALI
    fi
    
    dT=$(delta_time $tA)
    #--------------------------------
    #--- 2. POST-PROCESSING       ---
    #--------------------------------
    # create correct format
    #
    # CUSTOM OUTPUT:
    #       1      2      3      4      5     6        7      8       9       10   11     12   13   14     15
    #       qseqid sseqid evalue pident score bitscore length qcovhsp scovhsp qlen qstart qend slen sstart send"
    #
    # TARGET FORMAT:
    #    query   target  score evalue pid startQ endQ startD endD
    #
    info_msg "diamond: 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 "diamond: score scale mode = $ALIGNMENT_SCORE_MODE"
    {
        print_header "$EXE_DIAMOND blastp -d $dbase -q $query -o $__ALIGNMENT_TMPALI $__ALIGNMENT_OPTS" "$startdate"
        echo "# diamond align"
        echo "# ALIGNMENT_SCORE_MODE = $ALIGNMENT_SCORE_MODE"
        common_fmt_header
        trescore=$(time_msec)
        #if (( ALIGNMENT_SCORE_MODE == 0 )); then
        #    # S' = 2*S*alnlen/(qlen+tlen)
        #    grep -v "^#" "$__ALIGNMENT_TMPALI"    | 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)
        ####grep -v "^#" "$__ALIGNMENT_TMPALI"    | 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

        if (( ALIGNMENT_SCORE_MODE == 0 )); then
            # no change to score
            awk -v X="$MIN_OVERLAP" '!/^#/ && ($15-$14) >= X && ($12-$11) >= X { print $1,$2,$5,$3,$4,$11,$12,$14,$15; }' "$__ALIGNMENT_TMPALI"
            
        elif (( ALIGNMENT_SCORE_MODE == 1 )); then
            # normalise score according to sum of lengths
            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_TMPALI"

        elif (( ALIGNMENT_SCORE_MODE == 2 )); then
            # normalise score according to min length
            awk -v X="$MIN_OVERLAP" '!/^#/ && ($15-$14) >= X && ($12-$11) >= X { print $1,$2,100.0*$5/($10 < $13 ? $10 : $13),$3,$4,$11,$12,$14,$15; }' "$__ALIGNMENT_TMPALI"
        else
            # default: no change
            awk -v X="$MIN_OVERLAP" '!/^#/ && ($15-$14) >= X && ($12-$11) >= X { print $1,$2,$5,$3,$4,$11,$12,$14,$15; }' "$__ALIGNMENT_TMPALI"
        fi
        
        #elif (( ALIGNMENT_SCORE_MODE == 3 )); then
        #    # S' = S
        #    grep -v "^#" "$__ALIGNMENT_TMPALI"    | awk '{ print $1,$2,$5,$3,$4,$11,$12,$14,$15; }'
        
        #elif (( ALIGNMENT_SCORE_MODE == 4 )); then
        #    # use bit score
        #    grep -v "^#" "$__ALIGNMENT_TMPALI"    | awk '{ print $1,$2,$6,$3,$4,$11,$12,$14,$15; }'
        #fi
        echo "#TDETAIL: tot = $dT ; align = $dTalign"
        print_footer "$dT"
        
    } > "$output"

    #--------------------------------
    #--- 3. CLEANUP               ---
    #--------------------------------
    chmod 644 "$output"
    info_msg "alignment file: $output"
    return 0
}

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