#!/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/>.
'''
Cluster file tools.
Pipeline usage: yes
'''
#-----------------------------------------------------------------
# MAIN
class ClusterElementBase:
    """
    Format expected:
    clid  key     type  [  Nedges  start   end     rawscore        normscore       evalue]
    """
    def __init__(self, elems , taxsep=None):
        self.key = None
        if len(elems) > 1:
            ###self.clid    = int(elems[0])
            self.key     = elems[1]
            self.cltype  = -1 if len(elems) < 2 else int(elems[2])
            self.species = None
            self.geneid  = None
            if taxsep is not None:
                e = self.key.split(taxsep)
                if len(e)>1:
                    self.species = e[0]
                    self.geneid  = e[1]

#
    def __cmp__(self,other):
        if (self.key is None) or (not other.key is None):
            return -1
        if self.key>other.key:
            return 1
        elif self.key==other.key:
            return 0
        else:
            return -1

    def __repr__(self):
        return '{} {}'.format( self.key, self.cltype )

    def __str__(self):
        return self.__repr__()

    def good(self):
        return self.key is not None

class ClusterElement:
    """
    Format expected:
    clid  key     type    Nedges  start   end     rawscore        normscore       evalue
    """
    def __init__(self, elems, taxsep=None):
        self.key = None
        if len(elems)==9:
            self.key  = elems[1]
            self.cltype = int(elems[2])
            self.nedges = int(elems[3])
            self.start = int(elems[4])
            self.end = int(elems[5])
            self.pid = float(elems[6])
            self.score = float(elems[7])
            self.evalue = float(elems[8])
            self.geneid = None
            self.species = None
            if taxsep is not None:
                e = self.key.split(taxsep)
                if len(e)>1:
                    self.species = e[0]
                    self.geneid  = e[1]

#
    def __cmp__(self,other):
        if (self.key is None) or (not other.key is None):
            return -1
        if self.clid>other.clid:
            return 1
        elif self.clid==other.clid:
            return 0
        else:
            return -1

    def __repr__(self):
        return ('%8d %20s %8d %8d %8d %8d %8.1f %8.4f %8.2e')%(
            self.clid,
            self.key,
            self.cltype,
            self.nedges,
            self.start,
            self.end,
            self.pid,
            self.score,
            self.evalue)

    def __str__(self):
        return self.__repr__()

    def good(self):
        return self.key is not None

class Cluster:
    '''
    Container for one cluster.
    data : maps species id to cluster element
    stats: an instance of ClusterElementStats , set externally
    '''
    def __init__(self,numid=None):
        self.clid = numid
        self.data = {}
        self.cltypes = set()
        self.nelems = 0
        
    def clear(self):
        self.clid = None
        self.data = {}
        self.cltypes.clear()
        self.nelems = 0
        
    def add(self, celem ):
        if not celem.good():
            return False
        if self.clid is None:
            self.clid = celem.clid
        if celem.species not in self.data:
            self.data[celem.species] = []
        self.data[celem.species].append(celem)
        self.nelems += 1
        self.cltypes.add(celem.cltype)
        return True

    def nspec(self):
        return len(self.data)
    
    def dataSpecies(self,sp):
        if sp not in self.data:
            return None
        return self.data[sp]
        
    def hasSpecies(self,sp):
        return sp in self.data


def get_next( fh, taxsep=':', celemType=ClusterElement ):
    clid=None
    done=False
    eof=False
    prevpos=0
    line=''
    lnr=0
    cluster = Cluster()
    while not done:
        lnr=lnr+1
        prevpos = fh.tell()
        line = fh.readline()
        if len(line)==0:
            # empty line will only occur at the end of the file
            eof=True
            done=True
            continue

        line = line.strip().rstrip('\n')
        if len(line)==0 or line[0]=='#':
            continue

        elems = line.split()
        clid = elems[0]

        if cluster.clid is None or cluster.clid == clid:
            celem = celemType( elems, taxsep )
            cluster.clid = clid
            cluster.add( celem )
            continue

        # rewind file handler
        fh.seek(prevpos)
        done = True

    return (eof,cluster)
