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

import os
import sys
import argparse
import glob
from timeit import default_timer as timer

class Alignment:
    def __init__(self,clid=None,data_line=None):
        self.set(clid,data_line)
        
    def set(self,clid,data_line):
        self.clear()
        self.clid = clid
        self._extract_data( data_line )
        
    def clear(self):
        self.clid = None
        self.cluster_element = None
        self.keyA = None
        self.keyB = None
        self.score = 0
        self.evalue = 0
        self.pid = 0
        
    def isgood(self):
        return self.cluster_element is not None

    def betterthan( self, other ):
        if not self.isgood():
            return not other.isgood()

        # use the same logic as in brhclus package
        # sort by:
        # 1. score
        # 2. evalue
        # 3. pid
        # 4. keyA
        # 5. keyB
        # if all are equal, the function returns False
        rval=False;
        if self.score == other.score:
            if self.evalue == other.evalue:
                if self.pid == other.pid:
                    if self.keyA == other.keyA:
                        # all equal - set the one with the highest key as 'best'
                        rval = ( self.keyB > other.keyB )
                    else:
                        rval = ( self.keyA > other.keyA )
                else:
                    #higher percentage id is better
                    rval = ( self.pid > other.pid )
            else:
                # lower e-value is better
                rval = ( self.evalue < other.evalue )
        else:
            # higher score is better
            rval = ( self.score > other.score )
        return rval

    def _extract_data( self, data_line ):
        if data_line is None:
            return
        # split data line
        # expected: query,target,bits,evalue,pident,qstart,qend,tstart,tend
        #             0      1     2     3      4      5     6     7     8
        elems = data_line.split()
        if len(elems) != 9:
            self.cluster_element = None
            return
        self.keyA   = elems[0]
        self.keyB   = elems[1]
        self.score  = float(elems[2])
        self.evalue = float(elems[3])
        self.pid    = float(elems[4])
        
        #254 MMULA:000437 3 1 42 251 100 66.67 7.03e-42
        # create a cluster element string
        # <clid> <query key> <cluster type> <nedges> <query start> <query stop> <pid> <score> <evalue>
        self.cluster_element = '{} {} {} {} {} {} {} {} {}'.format(self.clid, self.keyA, options.cltype, '1', elems[5], elems[6], elems[4], elems[2], elems[3] )

def get_clid( fali ):
    '''
    filename follows pattern   <dir>/<label>_<clid>at<nodeid>.align
    '''
    label = os.path.basename( fali )
    d = label.split('_')[-1]
    p = d.find('at')
    if p<1:
        return None
    return d[0:p]

def read_alignments( falignments ):
    '''
    read alignments from all alignment files
    return as a dictionary
    result[key] = best alignment found
    '''
    result = {}
    nbadfiles = 0
    nbad = 0
    nok = 0
    nfiles = len(falignments)
    ifile = 0
    iperc = 0
    ipercLast = 0
    start = timer()
    # loop over all alignment files
    for fali in falignments:
        ifile += 1
        clid = get_clid( fali )
        if clid is None:
            nbadfiles += 1
            continue

        # open file and read each line
        with open(fali) as fh:
            for line in fh:
                line = line.strip().rstrip('\n')

                # skip commented or empty lines
                if len(line)==0 or line[0] in '#':
                    continue
                aliobj = Alignment( clid, line )
                
                # skip bad data lines
                if not aliobj.isgood():
                    nbad += 1
                    continue
                nok += 1

                # we have an OK data line
                if aliobj.keyA in result:
                    if aliobj.betterthan( result[aliobj.keyA] ):
                        result[aliobj.keyA] = aliobj
                else:
                    result[aliobj.keyA] = aliobj
                    
        iperc = int(ifile*100/nfiles)
        if iperc > ipercLast:
            stop = timer()
            #state = '{:>8}/{:>8}'.format(ifile,nfiles)
            state = '{:>3}%'.format(iperc)
            ipercLast = iperc
            print('{} total nr of keys read = {:<8}, time[s] = {:.3f}'.format(state,len(result),stop-start))
            start = timer()
        
        
    return result

def get_cluster_elements( alignments ):
    result = {}
    for key,ali in alignments.items():
        clid = ali.clid
        if clid not in result:
            result[clid] = []
        result[clid].append(ali.cluster_element)
    return result
    
def process( options ):
    # cluster file
    #fcluster = os.path.join( options.srcdir, '{}.og'.format(options.nodeid) )
    
    # profile pairwise directory
    # this directory contains files with names matching
    #   <label>_<clid>at<nodeid>.align
    dprofpwc = os.path.join( options.srcdir, options.nodeid, 'PWC', options.taxid )

    # get all alignments
    alipattern = '{}/{}_*at{}.align'.format(dprofpwc,options.taxid,options.nodeid)
    falignments = glob.glob(alipattern)
    if len(falignments) == 0:
        print('error: no alignments found matching pattern {}'.format(alipattern))
        return 255

    # dictionary of all alignments
    # key   == the first key of the alignement line
    # value == Alignment object
    alignments = read_alignments( falignments )

    with open( options.output, 'w' ) as fh:
        if len(alignments) == 0:
            # perhaps not an error...
            fh.write('#EMPTY\n')
            return 0

        # get all mapped cluster elements
        celements = get_cluster_elements( alignments )
        if len(celements) == 0:
            fh.write('#EMPTY\n')
        else:
            for clid,celems in celements.items():
                for c in celems:
                    fh.write(c+'\n')
        return 0
    return 255


parser = argparse.ArgumentParser(description='merge alignment files from ogmmprof_score into clusters')
parser.add_argument('-n', '--nodeid', dest='nodeid', help='node id of cluster file',    metavar='id', required=True)
parser.add_argument('-t', '--taxid',  dest='taxid',  help='taxid label of aligned user file', metavar='id', required=True)
parser.add_argument('-s', '--srcdir', dest='srcdir', help='directory containing cluster and profile alignments of given node id', metavar='dir', required=True)
parser.add_argument('-o', '--output', dest='output', help='output file', metavar='file', required=True)
parser.add_argument('-c', '--cltype', dest='cltype', help='cluster type for mapped gene ids', type=int, default=10, metavar='cltype', required=False)


options = parser.parse_args()

process( options )
