#!/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/>.
'''
This tool reads the time progress information from the brhclus stdout and analyses the data.
Pipeline usage: -
'''
from __future__ import print_function

import sys
import matplotlib.pyplot as plt
import numpy as np
import time

#
# Read the log file.
# Searches for lines starting with a time stamp followed by a '%'.
# The expected format is:
#
# <time stamp> % <L1>,<L1ok>,<L2>,<L2ok>,<L3>,<L3ok>,<L4>,<L4ok>,<T1>,<T1ok>,<T2>,<T2ok>,<T3>,<T3ok>,<T4>,<T4ok>,<N(clus)>,<modularity>
#
def read_data( file ):

    fh = open(file)
    first=True
    tvalprev=0
    tval0=0
    ind=0

    T=[]
    dt=[]
    pval=[]

    nl1=[]
    nl2=[]
    nl3=[]
    nl4=[]
    nl1ok=[]
    nl2ok=[]
    nl3ok=[]
    nl4ok=[]

    nt1=[]
    nt2=[]
    nt3=[]
    nt4=[]
    nt1ok=[]
    nt2ok=[]
    nt3ok=[]
    nt4ok=[]

    nclus=[]
    modularity=[]

    tval0=0
    for line in fh.readlines():
        line = line.strip().rstrip('\n')
        if len(line)>0 and line[0] not in '#':
            elems = line.split()
            if not ( len(elems)==3 or len(elems)==4 ): continue
            perc = elems[2]
            if not perc[-1]=='%': continue
            pval.append(int(perc[0:-1]))

            if len(elems)==4:
                ntstr = elems[3].split(',')
                nl1.append(int(ntstr[0]))
                nl1ok.append(int(ntstr[1]))
                nl2.append(int(ntstr[2]))
                nl2ok.append(int(ntstr[3]))
                nl3.append(int(ntstr[4]))
                nl3ok.append(int(ntstr[5]))
                nl4.append(int(ntstr[6]))
                nl4ok.append(int(ntstr[7]))

                nt1.append(int(ntstr[9]))
                nt1ok.append(int(ntstr[9]))
                nt2.append(int(ntstr[10]))
                nt2ok.append(int(ntstr[11]))
                nt3.append(int(ntstr[12]))
                nt3ok.append(int(ntstr[13]))
                nt4.append(int(ntstr[14]))
                nt4ok.append(int(ntstr[15]))

                nclus.append(int(ntstr[16]))
                modularity.append(float(ntstr[17]))

            tstr = elems[0]+' '+elems[1]
            tval = int(time.mktime(time.strptime(tstr, '%Y-%m-%d %H:%M:%S'))) - time.timezone
            if first:
                dt.append(0)
                T.append(0)
                tval0=tval
            else:
                dt.append(tval-tvalprev)
                T.append(tval-tval0)
            first=False
            tvalprev=tval
            ind=ind+1

    T=np.array(T)
    dt=np.array(dt)
    pval=np.array(pval)
    nt1=np.array(nt1)
    nt2=np.array(nt2)
    nt3=np.array(nt3)
    nt4=np.array(nt4)
    nt1ok=np.array(nt1ok)
    nt2ok=np.array(nt2ok)
    nt3ok=np.array(nt3ok)
    nt4ok=np.array(nt4ok)

    nl1=np.array(nl1)
    nl2=np.array(nl2)
    nl3=np.array(nl3)
    nl4=np.array(nl4)
    nl1ok=np.array(nl1ok)
    nl2ok=np.array(nl2ok)
    nl3ok=np.array(nl3ok)
    nl4ok=np.array(nl4ok)
    nclus=np.array(nclus)
    modularity=np.array(modularity)

    return (T,dt,pval,nl1,nl1ok,nl2,nl2ok,nl3,nl3ok,nl4,nl4ok,nt1,nt1ok,nt2,nt2ok,nt3,nt3ok,nt4,nt4ok,nclus,modularity)


def plot_data( xdata,T,dt,pval,nt1,nt1ok,nt2,nt2ok,nt3,nt3ok,nt4,nt4ok,nclus,modularity, xlabel ):
    fig=plt.figure()
    p11=fig.add_subplot(3,1,1)
    p21=fig.add_subplot(3,1,2, sharex=p11)
    p31=fig.add_subplot(3,1,3, sharex=p11)
#    p41=fig.add_subplot(4,1,4, sharex=p11)

    # plot timing
    p11.plot(xdata,dt)
    p11.grid()
    p11.set_ylabel('dT (secs)')

    X = xdata[1:]
    dnt1=np.diff(nt1)
    dnt1ok=np.diff(nt1ok)
    dnt2=np.diff(nt2)
    dnt2ok=np.diff(nt2ok)
    dnt3=np.diff(nt3)
    dnt3ok=np.diff(nt3ok)
    dnt4=np.diff(nt4)
    dnt4ok=np.diff(nt4ok)

    # plot cluster evolution of various cases
    p21.plot(X,dnt1,   label='TRI1')
    p21.plot(X,dnt1ok, label='TRI1 ok')
    p21.plot(X,dnt2,   label='TRI2')
    p21.plot(X,dnt2ok, label='TRI2 ok')
    p21.plot(X,dnt3,   label='TRI3')
    p21.plot(X,dnt3ok, label='TRI3 ok')
    p21.plot(X,dnt4,   label='TRI4')
    p21.plot(X,dnt4ok, label='TRI4 ok')
    p21.grid()
    p21.legend(loc='upper left')
    p21.set_ylabel('dN (counts)')

    # plot fraction success
    ant1=np.array(dnt1)*1.0
    ant1ok=np.array(dnt1ok)*1.0
    z=(ant1==0)
    ant1[z]=1.0
    ant1ok[z]=1.0
    rnt1=ant1ok/ant1

    ant2=np.array(dnt2)*1.0
    ant2ok=np.array(dnt2ok)*1.0
    z=(ant2==0)
    ant2[z]=1.0
    ant2ok[z]=1.0
    rnt2=ant2ok/ant2

    ant3=np.array(dnt3)*1.0
    ant3ok=np.array(dnt3ok)*1.0
    z=(ant3==0)
    ant3[z]=1.0
    ant3ok[z]=1.0
    rnt3=ant3ok/ant3

    ant4=np.array(dnt4)*1.0
    ant4ok=np.array(dnt4ok)*1.0
    z=(ant4==0)
    ant4[z]=1.0
    ant4ok[z]=1.0
    rnt4=ant4ok/ant4

    p31.plot(X,rnt1,label='TRI1')
    p31.plot(X,rnt2,label='TRI2')
    p31.plot(X,rnt3,label='TRI3')
    p31.plot(X,rnt4,label='TRI4')
    p31.grid()
    p31.legend(loc='upper right')
    p31.set_ylabel('fraction success')
    p31.set_xlabel(xlabel, fontsize=14)

#    p41.plot(xdata,modularity)
#    p41.set_ylabel('modularity')
#    p41.grid()
    
#
# START HERE
#

if len(sys.argv)<2:
    print('usage: ',sys.argv[0],' <input file>')
    sys.exit(1)

fname=sys.argv[1]
print('Processing file: ',fname)
(T,dt,pval,nl1,nl1ok,nl2,nl2ok,nl3,nl3ok,nl4,nl4ok,nt1,nt1ok,nt2,nt2ok,nt3,nt3ok,nt4,nt4ok,nclus,modularity)=read_data(fname)

plot_data( pval,T,dt,pval,nt1,nt1ok,nt2,nt2ok,nt3,nt3ok,nt4,nt4ok,nclus,modularity, 'percentage done' )
###plot_data( T,   T,dt,pval,nt1,nt1ok,nt2,nt2ok,nt3,nt3ok,nt4,nt4ok,nclus,modularity, 'time(s)' )

# show all plots
plt.show()
