#!/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
#------------------------------------------------
source $DIR_ORTHOLOGER_BIN/timing_base.sh

tm_alignment_X() {
    [ $# -eq 2 ] || return 255
    local A=$1
    local B=$2
    local lenA=$(tm_get_fastalen $A)
    local lenB=$(tm_get_fastalen $B)
    local nA=$(tm_get_fastanelem $A)
    local nB=$(tm_get_fastanelem $B)
    local X=$(calc "log(($lenA*$lenB/1000000.0)*($nA*$nB/1000.0))/log(10.0)")
    echo $X
}

tm_extract_time_alignment() {
    [[ $# -eq 2 ]] || return 255
    local A=$1
    local B=$2
    local f=$DIR_PROJECT/PWC/$A/${A}_${B}.$EXT_ALIGN
    [[ -f $f ]] || return 255
    tm_extract_time $f
}

# generate fit data for alignment
tm_alignment_init() {
    # check for existing data file
    local outf=$(tm_filename_data ALIGNMENT)
    [[ -f $outf ]] && return 0
    
    # nr of species
    local nspec=$PL_TODO_NLABELS
    # max nr of files
    local N2=$((nspec*nspec))
    local NLAST=0

    # randomly shuffle todo file, unless all are used
    local tf1=""
    tmp_make tf1 "tmgenali_f1"
    local tf2=""
    tmp_make tf2 "tmgenali_f2"
    if [[ $TM_MAX_NPOINTS -ge $N2 || $TM_MAX_NPOINTS -eq 0 ]]; then
        cp -f $PL_TODO $tf1
        cp -f $PL_TODO $tf2
        NLAST=$N2
    else
        shuf $PL_TODO > $tf1
        shuf $PL_TODO > $tf2
        NLAST=$TM_MAX_NPOINTS
    fi
    # loop
    local N=0
    local done=0
    verbose_info_msg "generating alignment timing data file"
    rm -f $outf

    local tmpout=""
    tmp_make tmpout "tmgenali_outp"
    local startdate=$(date)
    local tA=$(time_msec)
    local ntst=0
    for sA in $(cat $tf1); do
        # get length of fasta file A
        local lenA=$(tm_get_fastalen $sA)
        [[ ${#lenA} -gt 0 ]] || continue
        
        for sB in $(cat $tf2); do
            ntst=$((ntst+1))
            # get length of fasta file B
            local lenB=$(tm_get_fastalen $sB)
            [[ ${#lenB} -gt 0 ]] || continue
            
            # get alignment file
            f=$DIR_PROJECT/PWC/$sA/${sA}_${sB}.$EXT_ALIGN
            [[ -f $f ]] || continue

            # get data
            local T=$(tm_extract_time $f)
            [[ -n $T ]] || { error_msg "unable to extract time from alignment of $sA and $sB"; continue; }
            local node=$(tm_extract_machine $f)
            local X=$(tm_alignment_X $sA $sB)
            echo "$node $X $T" >> $tmpout

            # found, count it
            N=$((N+1))
            [[ $N -lt $NLAST ]] || { done=1; break; }
        done
        [[ $done -eq 0 ]] || break
    done
    info_msg "Found timing data for $N cases ($ntst processed)"
    
    local dT=$(delta_time $tA)
    add_header_footer $0 "$tmpout" "$outf" "$startdate" "$dT" "alignment timing with $N points"
}

tm_alignment_fit() {
    # generate data to fit
    tm_alignment_init || return 255
    
    # make sure file is there
    local timedata=$(tm_filename_data ALIGNMENT)
    [[ -f $timedata ]] || { info_msg "alignment timing data file does not exist ($timedata)"; return 0; }

    # make sure we have enough data points
    local npts=$(grep -v "#" $timedata | wc -l)
    [ $npts -gt 1 ] || { warn_msg "must have at least two points of timing data in order to fit"; return 255; }
    
    # make sure file is there
    local fitdata=$(tm_filename_fit ALIGNMENT)
    [[ -f $fitdata ]] && return 0 #{ error_msg "alignment timing fit file exists ($fitdata)"; return 255; }

    # get all nodes used
    local machines=( ALL $(cut -d " " -f 1 "$timedata" | sort -u | grep -v "^#") )
    [ ${#machines} -gt 0 ] || { error_msg "empty list of machines from timing data file <$timedata>"; return 255; }

    # tmp file for data subsets
    tmpdata=""
    tmp_make tmpdata "timedata" || exit 255

    # create output
    touch "$fitdata"
    
    # loop over all nodes
    local node=""
    for node in ${machines[@]}; do
        data=""
        # get data for node
        if [[ "$node" == "ALL" ]]; then
            data="$timedata"
        else
            grep "^$node" "$timedata" > $tmpdata
            data="$tmpdata"
        fi
        
        # fit data log(T) vs log(X), not X is already logged hence no "--tlogx"
        local result=""
        result=$($DIR_ORTHOLOGER_BIN/plot_xy.py --noplot --linfit -f "$data" -X 1 -Y 2 --tlogy --fitres)
        [ $? -gt 0 ] || echo "$node $result" >> "$fitdata"
    done
    return 0
}

tm_alignment_predict_all() {
    [ $# -eq 2 ] || return 0
    local fitdata="$1"
    local X="$2"
    
    # loop over fitted results
    local Tsum=0
    local Tmin=0
    local Tmax=0
    local N=0
    while read line; do
        [[ "${line:1}" == "#" ]] && continue
        data=( $line )
        [ ${#data[@]} -eq 5 ] || continue
        local node=${data[0]}
        ####[[ $node == "ALL" ]] && continue # skip 'all'
        local A=${data[1]}
        local B=${data[2]}
        local T=$(tm_predict $X $A $B)
        Tsum=$( calc "${Tsum}+$T" )
        N=$((N+1))
        cmptmin=$(calc "$T < $Tmin")
        cmptmax=$(calc "$T >= $Tmax")
        [[ $N -eq 1 || "$cmptmin" == "1" ]] && Tmin=$T
        [[ $N -eq 1 || "$cmptmax" == "1" ]] && Tmax=$T
        echo "$node $T"
    done < "$fitdata"
    local Tave=$(calc "$Tsum/$N")
    echo "STATS $Tave $Tmin $Tmax $N"
}

tm_alignment_predict() {
    [ $# -ge 1 ] || return 0
    local X=$1
    local node=""
    [ $# -ge 2 ] && node=$2
    
    # make sure file is there
    local fitdata=$(tm_filename_fit ALIGNMENT)
    [[ -f $fitdata ]] || return 0 #{ error_msg "alignment timing fit file does not exist ($fitdata)"; return 255; }

    [[ ${#node} -gt 0 ]] || node="ALL"
    data=( $(grep "^${node} " "$fitdata") )
    [ ${#data[@]} -eq 5 ] || return 255
    local A=${data[1]}
    local B=${data[2]}
    local T=$(tm_predict $X $A $B)
    echo "$T $T $T 1"

    return 0
}
