#!/usr/bin/python3
import sys
import argparse
import os

def read_nlast_line( filename, n, tag ):
    if n<1:
        n=1
    last_line = None
    nn=0
    try:
        with open(filename, 'rb') as f:
            try:  # catch OSError in case of a one line file 
                f.seek(-2, os.SEEK_END)
                done = False
                while not done:
                    if f.read(1) == b'\n':
                        nn=nn+1
                    if nn==n:
                        done=True
                    else:
                        f.seek(-2, os.SEEK_CUR)

            except OSError:
                pass

            for line in f:
                last_line = line.decode()
                if last_line.startswith(tag):
                    break
    except:
        pass

    return last_line

def find_at_end( filename, tag ):
    if not os.path.exists(filename):
        return None
    
    last_line = None    
    with open(filename, 'rb') as f:
        try:  # catch OSError in case of a one line file
            # seek back 128 bytes
            f.seek(-128, os.SEEK_END)
        except OSError:
            f.seek(0)

        for line in f:
            last_line = line.decode()
            if last_line.startswith(tag):
                break

    return last_line

class FastaStats:
    def __init__(self,ntot,minlen,maxlen,avelen,medlen,totlen):
        self.ntot=ntot
        self.minlen=minlen
        self.maxlen=maxlen
        self.avelen=avelen
        self.medlen=medlen
        self.totlen=totlen
    
def read_fasta_stats( filename ):
    data = FastaStats(0,0,0,0,0,0)
    with open(filename,'r') as fh:
        for line in fh:
            if len(line)==0:
                continue
            elems = line.split(':')
            if not len(elems)==2:
                continue
            if   elems[0]=='#NTOT':
                data.ntot = float(elems[1])
            elif elems[0]=='#MINLEN':
                data.minlen = float(elems[1])
            elif elems[0]=='#MAXLEN':
                data.maxlen = float(elems[1])
            elif elems[0]=='#AVELEN':
                data.avelen = float(elems[1])
            elif elems[0]=='#MEDLEN':
                data.medlen = float(elems[1])
            elif elems[0]=='#TOTLEN':
                data.totlen = float(elems[1])

    if data.ntot==0:
        return None
    return data

def read_align_usage( filename ):
    time_line = find_at_end( filename, '#TIME' )
    if time_line is None:
        return '-',0,0
    elems = time_line.split(':')
    if not len(elems)==2:
        return '-',0,0
    tres = elems[1].split()
    if not len(tres)==2:
        return '-',0,0
    node = tres[0]
    T = float(tres[1])
    fsize = os.path.getsize(filename)
    return node,T,fsize

def read_labels( filename ):
    labels=[]
    with open(filename,'r') as fh:
        for line in fh:
            if len(line)==0:
                continue
            if line[0]=='#':
                continue
            labels.append(line.rstrip())
    if len(labels)==0:
        return None
    return labels


def get_pwc_usage( options, A, B, ext ):
    f = os.path.join( options.pwc, '{}/{}_{}.{}'.format(A,A,B,ext))
    node,T,fsize = read_align_usage(f)
    return f,node,T,fsize



def process( options ):
    # read labels
    labels = read_labels( options.todo )
    if labels is None:
        print('error: missing or empty todo file at {}'.format(options.todo))
        return

    # read fasta stats for each label
    fstats={}
    Nstats=0
    for flbl in labels:
        f = os.path.join( options.sequences, '{}.fs_selected.stats'.format(flbl))
        data = read_fasta_stats(f)
        fstats[flbl] = data
        if data is not None:
            Nstats += 1
    if Nstats==0:
        print('error: no fasta stats was found')
        return

    # number of labels (organisms)
    N=len(labels)
    npairsD=N*(N+1)/2  # with diagonal
    npairsH=N*(N-1)/2  # without diagonal

    meanSeqN=0.0
    meanSeqLen=0.0
    meanABLen=0.0
    
    meanTimeAli=0.0
    meanTimeBRH=0.0
    meanTimeInp=0.0
    meanSizeAli=0.0
    meanSizeBRH=0.0
    meanSizeInp=0.0

    
    # process pairwise files
    print('A B nA lenA nB lenB aliNode aliT aliSize brhNode brhT brhSize inpNode inpT inpSize')
    for i in range(0,N):
        A = labels[i]
        Astats = fstats.get(A,None)
        if Astats is None:
            print('error: no stats for label {}'.format(A))
            continue
        meanSeqN += Astats.ntot
        meanSeqLen += Astats.avelen
        for j in range(i,N):
            B = labels[j]
            Bstats = fstats.get(B,None)
            if Bstats is None:
                print('error: no stats for label {}'.format(B))
                continue

            meanABLen += Astats.avelen*Bstats.avelen
            
            # alignment files
            f,ali_node,ali_T,ali_fsize = get_pwc_usage(options,A,B,'align')
            if ali_node == '-':
                print('error: unable to parse resource usage for align file {}'.format(f))
                continue
            meanTimeAli += ali_T
            meanSizeAli += ali_fsize
            
            # BRH files, not for i==j
            f,brh_node,brh_T,brh_fsize = get_pwc_usage(options,A,B,'brh')
            meanTimeBRH += brh_T
            meanSizeBRH += brh_fsize
            
            # inpar files, not for i==j
            f,inp_node,inp_T,inp_fsize = get_pwc_usage(options,A,B,'inpar')
            meanTimeInp += inp_T
            meanSizeInp += inp_fsize
            
            print(A,B,
                  Astats.ntot,Astats.avelen,
                  Bstats.ntot,Bstats.avelen,
                  ali_node,ali_T,ali_fsize,
                  brh_node,brh_T,brh_fsize,
                  inp_node,inp_T,inp_fsize )

    meanSeqN = meanSeqN/N
    meanSeqLen = meanSeqLen/N
    meanABLen  = meanABLen/npairsD
    meanTimeAli = meanTimeAli/npairsD
    meanSizeAli = meanSizeAli/npairsD
    meanTimeBRH = meanTimeBRH/npairsH
    meanSizeBRH = meanSizeBRH/npairsH
    meanTimeInp = meanTimeInp/npairsH
    meanSizeInp = meanSizeInp/npairsH

    print('#NLBLS        ',N)
    print('#NPAIRSD      ',npairsD)
    print('#NPAIRSH      ',npairsH)
    print('#MEANSEQN     ',meanSeqN)
    print('#MEANSEQLEN   ',meanSeqLen)
    print('#MEANABLEN    ',meanABLen)
    print('#MEANTIMEALI  ',meanTimeAli)
    print('#MEANSIZEALI  ',meanSizeAli)
    print('#MEANTIMEBRH  ',meanTimeBRH)
    print('#MEANSIZEBRH  ',meanSizeBRH)
    print('#MEANTIMEINP  ',meanTimeInp)
    print('#MEANSIZEINP  ',meanSizeInp)
    
parser = argparse.ArgumentParser(description='Summarise run resource usage')

parser.add_argument('-t', '--todo', dest='todo',
                    help='todo file', metavar='FILE', required=True)
parser.add_argument('-s', '--sequences', dest='sequences',
                    help='Sequences directory', metavar='string', required=False,default='Sequences')
parser.add_argument('-p', '--pwc', dest='pwc',
                    help='PWC directory', metavar='string', required=False,default='PWC')

options = parser.parse_args()

process(options)
