#!/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/>.
#
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# This is a special version of the alignment script that also extracts the BRHs.
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
#--- 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"

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

###nthreads=3

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

#---------------------------------
#--- 1. SETUP ALIGNMENT OPTS   ---
#---------------------------------
# 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
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
tmp_base parallel_tmp

# set opts
aliopts="$OPTS_DIAMOND --threads $nthreads -e $MAX_EVALUE --masking none --parallel-tmpdir $parallel_tmp"

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

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

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

# run aligner and output to tmpfile
tmp_query=""
tmp_make tmp_ali  diamond
tmp_make tmp_query  diamond

# remove comments
grep -v '^#' "$query" > "$tmp_query"
info_msg "diamond command:  $EXE_DIAMOND blastp -d ${dbase}.dmnd -q $query -o $tmp_ali $aliopts"
####<(grep -v '^#' "$query")
Talign=$(time_msec)
$EXE_DIAMOND blastp -d "${dbase}.dmnd" -q "$tmp_query" -o "$tmp_ali" $aliopts --quiet > /dev/null
ec=$?
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      : $aliopts"
        error_msg "# Output    : $SYS_UNAME:$tmp_ali"
        error_msg "# Error log : $(pwd)/${output}.ali.err"
    }
    #>& "${output}.ali.err"
    exit 255
fi

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

dT=$(delta_time $tA)
#--------------------------------
#--- 6. 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"

tmp_make tmp_ali_fmt  diamond
grep -v "^#" "$tmp_ali"    | awk '{ print $1,$2,100.0*$5/($10+$13),$3,$4,$11,$12,$14,$15; }' > $tmp_ali_fmt

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 $tmp_ali $aliopts" "$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 "^#" "$tmp_ali"    | 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)
    cat $tmp_ali_fmt
    grep -v "^#" "$tmp_ali"    | awk '{ print $1,$2,100.0*$5/($10+$13),$3,$4,$11,$12,$14,$15; }'
    
    #elif (( ALIGNMENT_SCORE_MODE == 3 )); then
    #    # S' = S
    #    grep -v "^#" "$tmp_ali"    | awk '{ print $1,$2,$5,$3,$4,$11,$12,$14,$15; }'
    
    #elif (( ALIGNMENT_SCORE_MODE == 4 )); then
    #    # use bit score
    #    grep -v "^#" "$tmp_ali"    | awk '{ print $1,$2,$6,$3,$4,$11,$12,$14,$15; }'
    #fi
    echo "#TDETAIL: tot = $dT ; align = $dTalign"
    print_footer "$dT"
    
} > "$output"

#--------------------------------
#--- 8. MAKE BRH              ---
#--------------------------------
# check findbrh extra options
extraopts=""
if [[ $MAKEBRH_NALT -gt 1 ]]; then
    extraopts="--Nalt $MAKEBRH_NALT --fsepmax $MAKEBRH_FSEPMAX"
fi
cmdargs="--AB \"$tmp_ali_fmt\" $extraopts"
{
    print_header "$EXE_FINDBRH $cmdargs"  "$startdate"
    echo "# ALIGNMENT_MATRIX = $ALIGNMENT_MATRIX"
    if (( MAKEBRH_NALT > 1 )); then
        echo "# fuzzy BRHs with N=$MAKEBRH_NALT and fsepmax=$MAKEBRH_FSEPMAX"
    fi
    common_fmt_header
    tA=$(time_msec)
    eval "$EXE_FINDBRH" "$cmdargs"
    ec=$?
    dT=$(delta_time $tA)
    print_footer "$dT"
}
#--------------------------------
#--- 8. CLEANUP               ---
#--------------------------------
chmod 644 "$output"
