#!/usr/bin/python3

# 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/>.
'''
Estimate the time usage for a given setup.
Experimental - use with care...
Pipeline usage: -
'''
from __future__ import print_function

import os
import sys
import math

def _secs2time(t):
    days = int(t/3600/24)
    t = t-days*3600*24
    hrs = int(t/3600)
    t = t-hrs*3600
    mins = int(t/60)
    secs = t - mins*60
    return ('%d days %2dh %2dm %2ds') % (days,hrs,mins,secs)

def get_elapsed_time(file):
    tact = 0
    f = open(file,'r')
    f.seek(-256,2)
    for line in f:
        if "# Time" in line:
            line = line.strip('\n')
            elems = line.split(':')
            tact = int(elems[1])
            break
    return tact

#-----------------------------------------------------------------
# MAIN

# check args
if len(sys.argv)<2:
    print("usage: "+sys.argv[0]+" <todo file> [sequence dir] [suffix]")
    sys.exit(1)

# set args
todo_file = sys.argv[1]
sequence_dir = 'Sequences'
suffix = '.fs97'
output_dir = 'PWC'
pair_suffix = '.sw_all'
include_same_pair = True

if len(sys.argv)>2:
    sequence_dir = sys.argv[2]
if len(sys.argv)>3:
    suffix = sys.argv[3]

# linear estimate
slope = 3250.0
offset = 5.0

# open todo file and read data
todoFH = open(todo_file)
lines = todoFH.readlines();
todoFH.close()
data = lines[0].split()

total_time = 0
tdone = 0
ndone = 0
ntodo = 0
nsdone = 0
nstodo = 0
ntot = 0
tmax=0
tdiffsum=0
tdiff2sum=0
# loop over all pairs and estimate time
for ia in range(0,len(data)):
    fnameA = sequence_dir+'/'+data[ia]+suffix
    szA = os.path.getsize(fnameA)

    for ib in range(0,len(data)):
        if (not include_same_pair) and (ib==ia):
            continue
        fnameB = sequence_dir+'/'+data[ib]+suffix
        pairname = output_dir+'/'+data[ia]+'_'+data[ib]+pair_suffix
        pairname2 = output_dir+'/'+data[ib]+'_'+data[ia]+pair_suffix
        szB = os.path.getsize(fnameB)
        nAB = (szA/10000.0)*(szB/10000.0)/1024.0/1024.0
        tAB = slope*nAB + offset
        ntot = ntot+1
        isdone = ''
        tactA=0
        tactB=0
        tact=0
        if os.path.exists(pairname):
            tactA = get_elapsed_time(pairname)
            tdone = tdone+tactA
            tdiff = tAB-tactA
            tdiffsum = tdiffsum + tdiff
            tdiff2sum = tdiff2sum + tdiff*tdiff
            nsdone = nsdone + nAB
            ndone = ndone+1
            isdone = '  DONE A'
        else:
            ntodo = ntodo + 1
            nstodo = nstodo + nAB
        
        if os.path.exists(pairname2):
            tactB = get_elapsed_time(pairname2)
            tdone = tdone+tactB
            tdiff = tAB-tactB
            tdiffsum = tdiffsum + tdiff
            tdiff2sum = tdiff2sum + tdiff*tdiff
            nsdone = nsdone + nAB
            ndone = ndone+1
            if len(isdone)>1:
                isdone = isdone + 'B'
            else:
                isdone = isdone + '  DONE  B'
        else:
            ntodo = ntodo + 1
            nstodo = nstodo + nAB

        if tactA>0 and tactB>0:
            tact = (tactA+tactB)/2
        else:
            if tactA>0:
                tact = tactA
            if tactB>0:
                tact = tactB

        if tact>0:
            tfract = tAB/tact
        else:
            tfract = -1
        
        if tAB>tmax:
            tmax = tAB
        total_time = total_time+tAB
        if ia != ib:
            total_time = total_time+tAB
            ntot = ntot+1

        if tfract<0:
            print('%s %s   :   %s' % ( data[ia], data[ib], _secs2time(tAB) ))
        else:
            print('%s %s   :   %s%s   %.1f%%' % ( data[ia], data[ib], _secs2time(tAB), isdone, tfract*100 ))
            
testim = nstodo*(tdone/nsdone)
print('-------------------------------------------------------')
print('Estimated total time      : ' + _secs2time(total_time))
#print 'Estimated remaining time  : %s (%5.1f%%)' % (_secs2time(total_time - tdone), 100*(total_time - tdone)/total_time )
print('Max estimated time        : ' + _secs2time(tmax))
print('Actual time done          : ' + _secs2time(tdone))
print('Estimated remaining time  : ' + _secs2time(testim))
print('Pairs done                : %d/%d (%5.1f%%)' % (ndone, ntot, 100.0*ndone/ntot))
tdiffsig=0
if ndone>1:
    tdiffsig = math.sqrt((tdiff2sum - (tdiffsum*tdiffsum/ndone))/ndone)
print('Mean difference est-true  : %5.2f (%5.2f) mins' % (tdiffsum/ndone/60.0, tdiffsig/60.0))
