#!/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
#-----------------------------------------------
# Global user settable variables
#-----------------------------------------------
# Variables used in check_pair_files()
# display empty results
CHK_SHOW_EMPTY=${CHK_SHOW_EMPTY:=1}
# check level as defined by check_file_cont()
# 0 : file existance
# 1 : + end tag
# 2 : + file size
CHK_LEVEL=${CHK_LEVEL:=1}
CHK_LEVEL_STR=("FAST" "ENDTAG" "FULL")

chk_msg_string() {
    [[ $# -eq 2 ]] || return 255
    local label=$1
    local level=$2
    local levelstr="X"
    [[ $level -le $CHK_LEVEL ]] && levelstr=${CHK_LEVEL_STR[$level]}
    
    printf "%-30s %-10s" "$label" "[$levelstr]"
}
#-----------------------------------------------
# check per file in Sequence
#-----------------------------------------------
check_per_file() {
    [[ $# -lt 3 ]] && fatal_msg "check_per_file() invalid nr of arguments = $#"
    CHECK_COUNTER=0
    local label="$1"
    local ext="$2"
    local chkwhat=$3
    local lblref=""
    local extref=""
    local fsz="0.01"
    [[ $# -ge 4 ]] && lblref="$4"
    [[ $# -ge 5 ]] && extref="$5"
    [[ $# -eq 6 ]] && fsz="$6"

    set_stepid $label

    local debug=0
    [[ $label == "SELECT" ]] && debug=1
    
    # get compression extension
    # make sure check level is correct
    # CHK_LEVEL is a global level
    [[ $chkwhat -le $CHK_LEVEL ]] || chkwhat=$CHK_LEVEL
    
    local nok=0
    local chkf=""
    local chkref=""
    local ec=0
    local i=0
    for ((i=0; i<${#FASTA_BASENAMES[@]}; i++)); do
        local fs="${FASTA_BASENAMES[$i]}"
        chkf="${fs}.${ext}"
        [[ -f $chkf ]] || { cnt_warn_msg "missing file $(trim_name $chkf)"; continue; }
        chkref=""
        if [[ ${#extref} -gt 0 ]]; then
            chkref="${fs//.${EXT_FASTA}/.${extref}}"
            [[ -f $chkref ]] || { cnt_warn_msg "missing file $(trim_name $chkref)"; continue; }
        fi
	    check_file_cont "$chkf" "$chkwhat" "$chkref" "$fsz"
	    ec=$?
	    [[ $ec -eq 0 ]] && nok=$((nok+1))
    done
    cnt_sum_msg
    local msg="${FACE_BOLD}$(chk_msg_string $label $chkwhat)${FACE_NORMAL}"
    num_status_msg "$msg" ${#FASTA_BASENAMES[@]} $nok
}

#-----------------------------------------------
#
#-----------------------------------------------
check_step_preproc() {
    # check existance plus for zero-size
    check_per_file "PREPROC" "${EXT_FASTA_CLEAN}" 1
}

check_step_masker() {
    # check existance, zero-size and relative size
    check_per_file "MASKER" "${EXT_FASTA_MASKED}" 2 "$(get_step_depend MASKER)" "${EXT_FASTA_CLEAN}" 0.99
}

check_step_select() {
    # check existance, zero-size and relative size
    check_per_file "SELECT" "${EXT_FASTA_SELECTED}" 2 "$(get_step_depend SELECT)" "${EXT_FASTA_MASKED}" 0.7
}

check_step_stats() {
    # check for existance, not zero size as the file contains only commented lines
    check_per_file "STATS" "${EXT_FASTA_STATS}" 0
}

check_step_formatdb() {
    CHECK_COUNTER=0
    local nok=0
    local chkf=""
    local ec=0
    local i=0
    for ((i=0; i<${#FASTA_BASENAMES[@]}; i++)); do
        local fs="${FASTA_BASENAMES[$i]}"
        chkf="${fs}.${EXT_FASTA_SELECTED}"
        check_formatdb_files "$chkf" 1
        ec=$?
	    [[ $ec -eq 0 ]] && nok=$((nok+1))
    done
    cnt_sum_msg
    local msg="${FACE_BOLD}$(chk_msg_string FORMATDB 0)${FACE_NORMAL}"
    num_status_msg "$msg" ${#FASTA_BASENAMES[@]} $nok
}
#-----------------------------------------------
# Check pair-wise files, fast mode.
#-----------------------------------------------
check_pair_files_fast() {
    [[ $# -eq 3 ]] || return 255
    local label="$1"
    local extension="$2"
    local matrix=$(get_matrix_type $label)

    local N=${#PL_TODO_LABELS[@]}
    local Ntot=0
    local T=""
    local nempty=0
    local nok=0

    local NB=$N
    case $matrix in
        0)  dNB=0  # full matrix
            ;;
        1)  dNB=1  # half-matrix, including diagonal
            NB=$((N+1))
            ;;
        2)  dNB=1  # half-matrix, excluding diagonal
            ;;
        3)  dNB=0  # full matrix, excluding diagonal
            NB=$((N-1))
            ;;
        *)  error_msg "invalid matrix type <$matrix>"
            return 255
            ;;
    esac

    # loop over all todo labels
    for T in "${PL_TODO_LABELS[@]}"; do

        # update nr of expected B
        NB=$((NB-dNB))
#        [[ $NB -gt 0 ]] || break
        Ntot=$((Ntot+NB))
        
        # check label
        isok_label "A" $T || continue

        # PWC directory - get this from function!!!
        local pwcdir="PWC/$T"
        [[ -d $pwcdir ]] || continue

        # get nr of files
        local nfiles=$(find "$pwcdir/" -maxdepth 1 -name '*.'$extension | wc -l)
        [[ $nfiles -eq 0 ]] && nempty=$((nempty+1))
        [[ $CHK_SHOW_EMPTY -eq 0 && $nfiles -eq 0 ]] && continue

        # print nr of files
        if [[ $VERBOSE -gt 0 ]]; then
            local msg="${FACE_BOLD}$(chk_msg_string $label 0) ${T}${FACE_NORMAL}"
            num_status_msg "$msg" $NB $nfiles
        fi
        [[ $NB -eq $nfiles ]] && nok=$((nok+NB))
    done
    cnt_sum_msg
    #
    if [[ $VERBOSE -eq 0 ]]; then
        local msg="${FACE_BOLD}$(chk_msg_string $label 0)${FACE_NORMAL}"
        num_status_msg "$msg" $Ntot $nok
    fi
}

#-----------------------------------------------
#
#-----------------------------------------------
check_pair_files() {
    [[ $# -lt 2 ]] && fatal_msg "check_pair_files() invalid nr of arguments = $#"
    if [[ $CHK_LEVEL -eq 0 ]]; then
        debug_msg "doing fast check of $1"
        check_pair_files_fast $1 $2
        return $?
    fi
    debug_msg "doing detailed check of $1"
    local label="${1}"   # step label
    local ext="${2}"     # file extension
###    local what="${3}"    # 0 -> full matrix, 1 -> half plus diagonal, 2 -> half except diagonal
    local chkwhat="0"    # what to check, 0 -> existance, 1 -> end tag, 2 -> size
    local lblref=""      # step id of reference
    local extref=""      # extension of reference
    local fsz="0.01"     # size fraction wrt reference
    [[ $# -ge 3 ]] && chkwhat="$3"
    [[ $# -ge 4 ]] && lblref="$4"
    [[ $# -ge 5 ]] && extref="$5"
    [[ $# -eq 6 ]] && fsz="$6"

    # make sure the check level is good
    [[ $chkwhat -le $CHK_LEVEL ]] || chkwhat=$CHK_LEVEL

    # get matrix type of step
    local matrix=$(get_matrix_type $label)
    
    # get index ranges
    # this returns the string:
    # <first A> <last A> <first B> <lastB>
    # Hence all indices are valid for the PL_TODO_LABELS array
    get_pairwise_matrix $matrix 0

    local beginA=${PAIRWISE_MATRIX[0]}
    local endA=${PAIRWISE_MATRIX[1]}
    local beginB=${PAIRWISE_MATRIX[2]}
    local endB=${PAIRWISE_MATRIX[3]}

    [[ $beginA -eq -1 ]] && { warn_msg "check_pair_files() got empty ranges in A,B for step $label"; return 0; }
    
    local numlabels=${#PL_TODO_LABELS[@]}
    local nokall=0
    local ntotall=0
    local msg=""

    local pairref=""
    local pairfile=""
    local pairbase=""
    local do_details=0
    local ec=0
    [[ $VERBOSE -eq 0 ]] || do_details=1
    
    CHECK_COUNTER=0

    # loop over label A
    local indA=0
    for (( indA=beginA; indA<endA+1; indA++ )); do
        local A="${PL_TODO_LABELS[$indA]}"
        # loop over label B
        local iBstart=$beginB
        local iBend=$endB
        
        # for half matrix, iBstart should be > indA
        [[ $matrix -eq 2 && $iBstart -le $indA ]] && iBstart=$((indA+1))
        [[ $matrix -eq 3 && $iBstart -le $indA ]] && iBstart=$indA

        # make sure iBstart remains sane
        [[ $iBstart -lt $beginB ]] && iBstart=$beginB
        
        local ntot=0
        local nok=0
        local indB=0
        for (( indB=iBstart; indB<iBend+1; indB++ )); do
            ntot=$((ntot+1))
            local B="${PL_TODO_LABELS[$indB]}"
            pairfile="$DIR_PROJECT_PWC/$A/${A}_${B}.$ext"
            if [[ ${#extref} -gt 0 ]]; then
                pairfile="$DIR_PROJECT_PWC/$A/${A}_${B}.$extref"
            fi
	        check_file_cont "$pairfile" "$chkwhat" "$pairref" "$fsz"
            ec=$?
	        [[ $ec -eq 0 ]] && nok=$((nok+1))
	    done

        ntotall=$((ntotall+ntot))
        nokall=$((nokall+nok))
        if [[ $do_details -ne 0 ]]; then
            if [[ $ntot -ge 0 ]]; then
                msg="${FACE_BOLD}$(chk_msg_string $label $chkwhat) ${A}${FACE_NORMAL}"
                num_status_msg "$msg" $ntot $nok
            fi
        fi
        CHECK_COUNTER=0
        verbose_info_msg "done check for $A"
    done
    cnt_sum_msg
    
    CHECK_COUNTER=0
    if [[ $VERBOSE -eq 0 ]]; then
        msg="${FACE_BOLD}$(chk_msg_string $label $chkwhat)${FACE_NORMAL}"
        num_status_msg "$msg" $ntotall $nokall
    fi
}

#
# check alignment results
# - .align files
# should exist for all [A,B] and [B,A]
# check for zero sized files
#
check_step_alignment() {
    check_pair_files "ALIGNMENT" "$EXT_ALIGN" 1
    return $?
}

#
# check BRH files
# - per species : PWC/*.brh
# - total file : Cluster/*.set-brh
#
check_step_makebrh() {
    check_pair_files "MAKEBRH" "$EXT_BRH" 1 "ALIGNMENT" "$EXT_ALIGN" 0.001
    return $?
}

#
# check makeinparsel step
#
check_step_makeinparsel() {
    check_per_file MAKEINPARSEL $EXT_INPAR_SELECTED 1
    return $?
}

#
# check inparalogs
# - per species PWC/*.brh_inparalogs
# - all species Cluster/*.inparalogs
# check only for existance!
#
check_step_makeinpar() {
    check_pair_files "MAKEINPAR" "$EXT_INPAR" 1
    return $?
}

#
# Check file in Cluster directory
#
check_cluster_file() {
    [[ $# -lt 2 ]] && fatal_msg "check_cluster_file() invalid nr of arguments = $#"
    local stepid="$1"
    local ext="$2"
    local lblref=""
    local extref=""
    local fsz=0.99999
    [[ $# -ge 3 ]] && lblref="$3"
    [[ $# -ge 4 ]] && extref="$4"
    [[ $# -eq 5 ]] && fsz="$5"

    CHECK_COUNTER=0
    local testOK=0
    local fcl=""
    local nerr=0
    fcl="${DIR_PROJECT_CLUSTER}/${PL_TODO_NAME}.$ext"
    [[ -f $fcl ]] || { cnt_error_msg "missing file $fcl"; nerr=1; }
    if [[ $nerr -eq 0 ]]; then
        local fref=""
        if [[ ${#extref} -gt 0 ]]; then
            fref="${DIR_PROJECT_CLUSTER}/${PL_TODO_NAME}.$extref"
            [[ -f $fcl ]] || { cnt_error_msg "missing file $fref"; nerr=1; }
        fi

        # check file for
        # - existance
        # - zero size
        # - small size relative full BRH file
        check_file_cont "$fcl" $CHK_LEVEL "$fref" "$fsz"0
        [[ $? -ne 0 ]] || testOK=1
    
        cnt_sum_msg

    fi
    local nok=1
    [[ $testOK -ne 0 ]] || nok=0

    # Ntot should be always 1 unless no todo file is defined
    local ntot=1
    [[ -f "$PL_TODO" ]] || ntot=0

    local msg="${FACE_BOLD}$(chk_msg_string $stepid $CHK_LEVEL)${FACE_NORMAL}"
    num_status_msg "$msg" $ntot $nok
}

#
# Check final data sets - input to clustering
#
check_step_setbrh() {
    check_cluster_file "SETBRH" "$EXT_BRH"
}
check_step_setinpar() {
    check_cluster_file "SETINPAR" "$EXT_INPAR"
}
check_step_setinparsel() {
    check_cluster_file "SETINPARSEL" "$EXT_INPAR_SELECTED"
}


#
# check clustering
# - output from brhclus: Cluster/*.og_raw
#
check_step_cluster() {
    check_cluster_file "CLUSTER" "$EXT_CLUSTER_RAW"
}

#
# check merging of inparalogs
# - main inparalogs merged in Cluster/*.og_inpar
#
check_step_mergeinpar() {
    check_cluster_file "MERGEINPAR" "$EXT_CLUSTER_INPAR" "CLUSTER" "$EXT_CLUSTER_RAW" "0.9999"
}

#
# check merging of 97% inparalogs
# - final (including 97% id) in Cluster/*.og
#
check_step_mergeinparsel() {
    check_cluster_file "MERGEINPARSEL" "$EXT_CLUSTER" "MERGEINPAR" "$EXT_CLUSTER_INPAR" "0.9999"
}

#
# check post-processing
#
check_step_postproc() {
    [[ ${OP_STEP_ENABLE[POSTPROC]} -eq 1 ]] || {
        disabled_status_msg "${FACE_BOLD}POSTPROC${FACE_NORMAL}"
        return 0
    }
    # loop over all scripts
    local ec=0
    local s=
    for s in $OP_POSTPROC_SCRIPTS_FOUND; do
        local lbl=$(get_postproc_name "$s")
        # set postproc script
        set_postproc $s
        ec=$?
        if [[ $ec -ne 0 ]]; then
            error_status_msg "${FACE_BOLD}POSTPROC:$lbl${FACE_NORMAL}"
            continue
        fi

        # check all targets
        postproc_conf
        local nok=0
        local ntot=0
        local p=
        for p in $(echo $POSTPROC_TRGS); do
	        ntot=$((ntot+1))
	        check_file_cont "$p" $CHK_LEVEL
	        [[ $? -ne 0 ]] || nok=$((nok+1))
        done
        CHECK_COUNTER=0
        cnt_sum_msg
        local msg="${FACE_BOLD}$(chk_msg_string POSTPROC:$lbl $CHK_LEVEL)${FACE_NORMAL}"
        num_status_msg "$msg" $ntot $nok

        # reset postproc variables
        reset_postproc
    done
}
