#!/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
#------------------------------------------------
# define nr of targets - set in postproc_conf
export POSTPROC_NTARGETS=0
export POSTPROC_TRGS=("")

postproc_help() {
    echo -e \
	" ${str_help_input}\n" \
	" * PWC/<spec>/<spec>_*.align  - all alignment files\n" \
	" * Cluster/<todo>.og          - final clustering\n" \
	"${str_help_output}\n" \
	" * PWC/<spec>/<spec>_<todo>.relog_part  - partial result for <spec>\n" \
	" * Cluster/<todo>.relog                 - the final related group file\n" \
	"${str_help_desc}\n" \
	"This post-processing plugin calculates related group metrics.\n" \
    "In the tree mode it will only compute for the root node.\n" \
	"The first step - calculating the partial results - is time consuming.\n" \
	"For large data sets, consider using many threads (OP_STEP_NTHREADS[POSTPROC]).\n" \
    "Note that the calculation can be interrupted and resumed.\n" \
	"\n" \
	"The alignments are filtered on the e-value with a max allowed MAX_EVALUE.\n" \
	"If the data set is large, consider lowering this value before running.\n" \
	"That will effectively reduce the nr of possible links that goes into the calculation.\n" \
    "\n" \
	""
}

postproc_conf() {
    reset_joblists

    if (( TREE_ENABLED == 1 )); then
        [[ -f $TREE_INPUT_FILE ]] || step_fatal_msg "missing tree file $TREE_INPUT_FILE"
        get_internal_nodes || {
            step_error_msg "postproc_conf: unable to get nodes from $TREE_INPUT_FILE"
            return 255
        }
        for node in ${_NODES[@]}; do
            local ftrg="$DIR_PROJECT_CLUSTER/${node}.relog"
            POSTPROC_TRGS+=( "$ftrg" )
            JOBSUBLBL+=( "$JOBLOGNOSPEC" ) # not sure this is needed
        done
    else
        JOBSUBLBL=( "$JOBLOGNOSPEC" )
        POSTPROC_TRGS=("$DIR_PROJECT_CLUSTER/${PL_TODO_NAME}.relog")
    fi
    # get nr of targets - used by step_start in step_postproc() function
    POSTPROC_NTARGETS=${#POSTPROC_TRGS[@]}
}

postproc_cleanup() {
    local indA=0
    for (( indA=0; indA<$nspec; indA++ )); do
        local A=${PL_TODO_LABELS[$indA]}

        # check for required species
        has_required_label $A $A
        [[ $? -ne 0 ]] && continue

        # get pair dir for PWC A
        get_pwc_dir "$A"

        # target file
        local relogpart=$R_PWC_DIR/${A}_${PL_TODO_NAME}.relog_part
        rmcmd "$relogpart"
    done
    local f=""
    for f in ${POSTPROC_TRGS[@]}; do
        rmcmd "$f"
    done
}

#-------------------------------------------------------------------------------
create_partials() {
    echo "+++++ NOT YET IMPLEMENTED create_partials +++++"
    exit 255
    
    step_info_msg "create partial relog files for each node"
    is_dryrun && return 0

    
    # get all nodes
    [[ -f $TREE_INPUT_FILE ]] || step_fatal_msg "missing tree file $TREE_INPUT_FILE"
    get_internal_nodes || {
        step_error_msg "postproc_conf: unable to get nodes from $TREE_INPUT_FILE"
        return 255
    }
    step_info_msg "found ${#_NODES[@]} nodes"
    
    reset_joblists
    local node=""
    for node in ${_NODES[@]}; do
        local ftodo=$DIR_PROJECT_TODO/${node}.todo
        local fog=$DIR_PROJECT_CLUSTER/${node}.$EXT_CLUSTER
        local ftrg=$DIR_PROJECT_CLUSTER/${node}.$EXT_CLUSTER_OGXML
        # check if target should be remade
        remake_target "$ftrg" "$fog"
        [[ $? -eq 0 ]] || {
            step_info_msg "skipping $ftrg"
            continue
        }

        SRCFILES+=($ftodo)
        TRGFILES+=($ftrg)

        step_info_msg "adding $ftrg"
        
        # set JOB specs <root,node,nthreads>
        JOBARGS+=( "${fog},${ftrg},$ftodo" )
        JOBLABELS+=( "ogxml-$node" )
        JOBSUBLBL+=( "$JOBLOGNOSPEC" )
    done
    step_info_msg "start running each node (${#_NODES[@]} jobs)"
    if [[ ${#_NODES[@]} -eq 0 ]]; then
        step_info_msg "all nodes skipped - nothing to do"
        reset_joblists
        return 0
    fi
    
    # third argument flags that it should NOT call init_step() and step_done()
    step_per_label POSTPROC "${DIR_ORTHOLOGER_BIN}/run_ogxml.sh" "0"

    # check error logs
    if [[ ${#JOBLABELS[@]} -gt 0 ]]; then
        check_error $ec "failed ${STEP_ID}:$POSTPROC_NAME"
        [[ $? -eq 0 ]] || return 1
    fi

    # cleanup
    reset_joblists

}
#-------------------------------------------------------------------------------










postproc() {
    local rval=0

    # check if dry run
    is_dryrun && return 0

    local nparallel=$(get_step_nparallel)
    local nthreads=${OP_STEP_NTHREADS[$STEP_ID]}
    local isbatch=0
    is_runmode_batch && isbatch=1

    local nfailed=0
    local nsub=0
    local ec=0

    reset_joblists

    local nspec=$PL_TODO_NLABELS
    [[ $nspec -gt 0 ]] || { error_msg "no species defined by TODO ($PL_TODO)"; return 255; }

    # loop over first label
    # note that the last label is skipped (nspec-1)
    # -> using only half matrix without diagonal -> none to be done 
    local lastind=$((nspec-2))
    local indA=0
    local ndone=0
    for (( indA=0; indA<nspec-1; indA++ )); do
        local A=${PL_TODO_LABELS[$indA]}

        # check for required species
        has_required_label $A $A
        [[ $? -ne 0 ]] && continue

        # get pair dir for PWC A
        get_pwc_dir "$A"
        [[ -d $R_PWC_DIR ]] || fatal_msg "missing PWC dir for $A ($R_PWC_DIR)"

        # target file
        local relogpart=$R_PWC_DIR/${A}_${PL_TODO_NAME}.relog_part

        # check if relogpart needs remaking
        local remake=0

        # if target does not exist, need remake obviously
        [[ -f $relogpart ]] || remake=1

        # loop over all species B and make sure all input alignment files are there
        local indB=0
        local indB0=0
        [[ $ALIGNMENT_MATRIX -eq 0 ]] || indB0=$indA
        for (( indB=indB0; indB<nspec; indB++ )); do
            # get alignment file and die if it does not exist
            local B=${PL_TODO_LABELS[$indB]}
            local pairAB=${R_PWC_DIR}/${A}_${B}.$EXT_ALIGN
            [[ -f $pairAB ]] || pairAB=${pairAB}.gz # check gzipped
            [[ -f $pairAB ]] || fatal_msg "missing required alignment file $pairAB"

            # if remake condition is already set, skip to next
            # like this, all alignments are checked for existance
            [[ $remake -eq 0 ]] || continue

            # check if target should be remade
            remake_target "$relogpart" "$pairAB"
	        ec=$?
            [[ $ec -ne 0 ]] || remake=1
        done

        # nr of done
        ndone=$((ndone+1))
        
        # check if it should be made for species A
        if [[ $remake -ne 0 ]]; then
            # append to SRCFILES
            SRCFILES[${#SRCFILES[@]}]="$A" #TODO: should be actual source file
            JOBARGS[${#JOBARGS[@]}]="${A},${nthreads}"
            JOBLABELS[${#JOBLABELS[@]}]="${A}-${PL_TODO_NAME}"
            JOBSUBLBL[${#JOBSUBLBL[@]}]="$A"
        fi

        # submit all jobs if sufficiently many have been collected
        # if last A, force the submission
        local nsrc=${#SRCFILES[@]}

        # submit if conditions are met
        if [[ $nsrc -ge $nparallel ]] || [[ $indA -eq $lastind ]]; then
            # if no jobs, move to next A
            [[ $nsrc -gt 0 ]] || continue

	        nsub=$((nsub+nsrc))
            # check that the scheduler is available
            # wait if not
            check_scheduler_status $nsrc

            # submit jobs
            step_info_msg "<$A> submitting $nsrc related og jobs (per species $ndone/$nspec)"
            run_cmd "${DIR_ORTHOLOGER_BIN}/run_siblings.sh"
            ec=$?
            [[ $ec -eq 0 ]] || nfailed=$((nfailed+1))
            # if batch, then sleep a bit in order for the scheduler to calm down, then continue
            if [[ $isbatch -eq 1 ]]; then
                sj_wait || { step_error_msg "aborting step since job id wait failed"; exit 255; }
            fi

            # check error (execution of alignments have finished)
            check_error $ec "failed postproc_siblings"
            [[ $? -eq 0 ]] || return 1

            reset_joblists
        fi
    done

    step_info_msg "partial results done"
    
    # check if any src files are remaining
    # there may be remaining files if the last pair in loop above
    # was excluded 
    [[ ${#SRCFILES[@]} -eq 0 ]] || fatal_msg "postproc_siblings failed since srcfiles remain at the end of the loop (N= ${#SRCFILES[@]})"

    # clear job related arrays
    reset_joblists

    #------------------------------------
    # PART 2 merge all partial results
    #------------------------------------
    # if any of the partial job submissions failed, quit here
    [[ $nfailed -eq 0 ]] || return 1

    # target file
    local trgfile=$DIR_PROJECT_CLUSTER/${PL_TODO_NAME}.relog
    if [[ -f $trgfile ]]; then
	    # target exists and no partial ones were submitted, assume we are done
	    [[ $nsub -gt 0 ]] || {
            step_info_msg "seems like all is done"
            return 0
        }
    fi

    step_info_msg "merge partial results"
    SRCFILES[0]="1" #TODO: fix real name
    JOBARGS[0]="1"
    JOBLABELS[0]="merge-relog"
    JOBSUBLBL[0]="$JOBLOGNOSPEC"

    check_scheduler_status ${#JOBARGS[@]}

    # submit jobs
    step_info_msg "submitting ${#JOBARGS[@]} related og jobs (merge partial results)"
    run_cmd "${DIR_ORTHOLOGER_BIN}/merge_siblings.sh"
    ec=$?
    if [[ $ec -ne 0 ]]; then
        step_error_msg "failed to submit job with error code $ec"
        force_remake_new
        return 1
    fi

    # if batch, wait
    if [[ $isbatch -eq 1 ]]; then
        sj_wait || { step_error_msg "aborting step since job id wait failed"; exit 255; }
    fi

    # check error log
    check_error $ec "failed siblings partial step"
    [[ $? -eq 0 ]] || return 1

    link_results "_siblings.txt" $trgfile
    reset_joblists

    return 0
}
