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

SACCT_BIN=`which sacct 2>/dev/null`
[ -x "$SACCT_BIN" ] || { error_msg "slurm command <sacct> could not be found - is slurm installed?"; exit 255; }

declare -A MEMUSAGE
declare -A CPUUSAGE
declare -A LBLCOUNT

calc(){ awk "BEGIN{ print $* }" ;}

update_usage() {
    local label="$1"
    local mem="$2"
    local cpu="$3"
    #
    # First process memory info
    #
    # get last char of mem
    local msz="${mem:${#mem}-1:1}"
    # get mem value
    local memval=${mem::-1}
    [ ${#memval} -eq 0 ] && memval=0
    case $msz in
        K)  memval=`calc ${memval}*1024`
            ;;
        M)  memval=`calc ${memval}*1024*1024`
            ;;
        G)  memval=`calc ${memval}*1024*1024*1024`
            ;;
        *)  ;;
    esac

    if [ "${MEMUSAGE[$label]+1}" == "" ]; then
        MEMUSAGE[$label]=$memval
        LBLCOUNT[$label]=1
    else
        local oldmem=${MEMUSAGE[$label]}
        local count=${LBLCOUNT[$label]}
        MEMUSAGE[$label]=$((memval+oldmem))
        LBLCOUNT[$label]=$((count+1))
    fi
    #
    # Second, CPU info
    # Time may be given as
    # <days>-<hrs>:<min>:<secs>
    #
    # NOTE: values in T may have leading 0's and thus be interpreted as octal.
    # Hence '09' will cause an error. The '10#' enforces a decimal base.
    #
    declare -a T
    IFS=":-" read -ra T <<< "$cpu"
    local timeval=0
    if [ ${#T[@]} -eq 4 ]; then
        timeval=$((10#${T[0]}*24*3600 + 10#${T[1]}*3600 + 10#${T[2]}*60 + 10#${T[3]}))
    elif [ ${#T[@]} -eq 3 ]; then
        timeval=$((10#${T[0]}*3600 + 10#${T[1]}*60 + 10#${T[2]}))
    fi
    if [ "${CPUUSAGE[$label]+1}" == "" ]; then
        CPUUSAGE[$label]=$timeval
    else
        local oldtime=${CPUUSAGE[$label]}
        CPUUSAGE[$label]=$((timeval+oldtime))
    fi
}

getdata() {
    local sfile="$1"
    [ ! -f "$sfile" ] && return
    local showzero=$2

    local sbase=`basename $sfile`
    local jobid=`echo $sbase | sed 's%.*-\([0-9]*\)\.log%\1%'`
    local jobname=`echo $sbase | sed 's%\(.*\)-[0-9]*\.log%\1%'`
    local data=`$SACCT_BIN -j ${jobid} --format MaxRSS,MinCPU --noheader 2>/dev/null | tail -n 1`
    # remove trailing spaces - it will be empty if job is still running
    data=`echo "$data" | sed 's/ *$//g'`
    [ ${#data} -eq 0 ] && return

    read -ra ADATA <<< $data
    local td="${ADATA[0]}${ADATA[1]}"

    [ "$td" == "000:00:00" ] && [ $showzero -eq 0 ] && return

    update_usage "$jobname" "${ADATA[0]}" "${ADATA[1]}"
    printf "%-12s%-48s%-20s%s\n" "${jobid}" "${jobname}" "${ADATA[0]}" "${ADATA[1]}"
}

usage() {
    echo
    echo "usage: ${0} -d <directory with slurm files> [ -t <todo label> ] [ -s <stepid> ] [ -z ]"
    echo
    echo "       -z      include zero-sized jobs"
    echo "       -t      optional todo label"
    echo "       -s      optional step id"
}

joblogdir=
todolabel=""
stepid=""
showzero=0
while getopts hd:t:s:z OPT; do
    case "$OPT" in
	    h)	usage
            return 0
            ;;
        d)  joblogdir=$OPTARG
            ;;
        t)  todolabel=$OPTARG
            ;;
        s)  stepid=$OPTARG
            ;;
        z)  showzero=1
            ;;
        \?) fatal_msg "$0 - argument parse error"
            ;;
    esac
done


[ ${#joblogdir} -eq 0 ] && { usage; fatal_msg "missing -d <joblog dir> option!"; }

[ -d "$joblogdir" ] || fatal_msg "job log directory is not readable/existing ( $joblogdir )"



# get search pattern for job logs
sjbase=`sj_joblog_base_args_nopid $todolabel $stepid`
#[ ${#sjfullbase} -eq 0 ] || sjbase=`basename $sjfullbase`

echo find ${joblogdir}/ -maxdepth 1 -name "${sjbase}*.log"

# get all slurm joblogs
allslurms=($(find ${joblogdir}/ -maxdepth 1 -name "${sjbase}*.log" | sort))
if [ ${#allslurms[@]} -eq 0 ]; then
    echo "No jobs found"
    exit 0
fi

echo
echo "------------------------------------------------------------------------"
echo "----------------------R E S O U R C E S T A T S-------------------------"
echo "------------------------------------------------------------------------"
printf "%-12s%-48s%-20s%s\n" "Job ID" "Job Label" "Memory" "CPU time"
echo "------------------------------------------------------------------------"
for x in ${allslurms[@]}; do
    getdata $x $showzero
done

echo
echo "------------------------------------------------------------------------"
echo "---------------------------S U M M A R Y--------------------------------"
echo "------------------------------------------------------------------------"
printf "%-48s%15s%15s\n" "LABEL" "MEMORY (M)" "TIME (sec)"
echo "------------------------------------------------------------------------"
for i in "${!MEMUSAGE[@]}"; do
    memval=`awk "BEGIN{ print ${MEMUSAGE[$i]}/1024.0/1024.0/${LBLCOUNT[$i]} }"`
    cpuval=`awk "BEGIN{ print ${CPUUSAGE[$i]}/${LBLCOUNT[$i]} }"`
    printf "%-48s%15.1f%15.1f\n" "$i" $memval $cpuval
done
echo
