#!/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 code checks for a given fasta file if the sequences are proteins or nucletides.
#
import re
import orthologer_tools as OT

#-----------------------------------------------------------------
# LETTERS
# Note that 'N' is not in the protein list as it is ambiguous.
# 'U' is for RNA.
# For proteins there are a few letters with special meaning:
#  1. B - D or N
#  2. Z - E or Q
#  3. J - I or L
#
# Stop codons:
#  'U','O','*','-'
#

SeqLettersCommon    = ['A','T','C','G']
SeqLettersNucl      = []
SeqLettersProt      = ['O','R','D','Q','E','H','I','L','K','M','F','P','S','W','Y','V']
SeqLettersProtAmbig = ['B','Z','J']
SeqLettersMasked    = ['X','N']
SeqLettersStopCodon = ['*','-','O','U']
RegexCodons=r'UO\-*'

def getMasked(stype):
    char=None
    if stype=='NUCL':
        char='N'
    elif stype=='PROT':
        char='X'
    return char

def getSpecies( label, specsep ):
    t=label.split()
    if len(t)>1: return None
    # find split at ':'
    tokens=label.split(specsep)
    if not len(tokens)==2: return None
    return tokens[0]

def getSpeciesFile( fasta, specsep ):
    speclabel=None
    found=False
    while not found:
        # get next fasta entry
        (eof,label,descr,sequence)=getNext(fasta)
        # if sequence is empty means that an illegal line is found - stop!
        if sequence is None:
            break
        # skip line if empty or if it does not start with '>'
        if len(sequence)==0: continue
        speclabel=getSpecies(label, specsep )
        # stop if type is established
        if speclabel is not None: found=True

    fasta.seek(0,0)
    return found,speclabel

def isValidLabel( label, specsep ):
    if label is None: return False
    if specsep is None: return True
    # make sure there are no whitespace
    t=label.split()
    if len(t)>1: return False
    # find split at ':'
    tokens=label.split(specsep)
    if not len(tokens)==2: return False
    return True

def replaceLowerCaseMasked( fasta, lchr ):
    return re.sub('[a-z]', lchr, fasta)

def getNext( fasta ):
    description=None
    label=None
    sequence=None

    # search for line with '>' in the beginning
    foundName=False
    done=False
    eof=False
    prevpos=0
    myline=''
    lnr=0
    while not done:
        lnr=lnr+1
        prevpos = fasta.tell()
        myline = fasta.readline()
        if len(myline)==0:
            # empty line will only occur at the end of the file
            eof=True
            done=True
        elif not foundName:
            myline = myline.strip().rstrip('\n')
            if len(myline)==0: continue
            if myline[0]=='>':
                name = myline[1:]
                tokens=name.split()
                label=tokens[0]
                if len(tokens)>1: description=' '.join(tokens[1:])
                foundName = True
            elif myline[0]=='#':
                continue
            else:
                # found a non label line
                label = None
                sequence = None
                description = None
                done = True
        else:
            myline = myline.strip().rstrip('\n')
            if len(myline)==0: continue
            if myline[0]=='#': continue
            if myline[0]=='>':
                fasta.seek(prevpos)
                done = True
                continue
            if sequence is None: sequence=myline
            else: sequence = sequence+myline

#    if not foundName:
#        eof=1

    # remove possible '.' at end of sequence
    if sequence is not None:
        if sequence[-1]=='.':
            sequence = sequence[:-1]

    return (eof,label,description,sequence)

def parseResult(stype,eof,label,descr,sequence,specsep,labels,verbose):
    done=False
    skip=False
    errmsg=None

    if eof:
        done=True

    if sequence is None:
        return (done,True,'empty sequence at label <%s>'%label,None)

    # check label
    if specsep is not None:
        if not isValidLabel(label,specsep):
            errmsg='invalid label <%s>'%label
            skip=True

    # if label is not in list, skip
    if labels is not None:
        if label not in labels: skip=True

    # get fasta type
    if (stype is None) or (stype not in ['NUCL','PROT']):
        found,stype=getType(sequence,verbose)

    # bad fasta sequence found - reset type to None
    if stype=='BADT':
        stype=None
        errmsg='bad sequence for label <%s>'%label
        skip = True

    return (done,skip,errmsg,stype)

def readFasta(inpf,specsep,labels=None,verbose=False):

    # loop over all fasta
    done=False
    errmsg=None
    errcode=0
    stype=None
    species=None
    fastaDB={}
    speciesDB=[]

    while not done:
        # get next fasta
        (eof,label,descr,sequence)=getNext(inpf)

        # parse result
        (done,skip,errmsg,stype)=parseResult(stype,eof,label,descr,sequence,specsep,labels,verbose)
        if errmsg is not None:
            if not skip:
                errcode=1
                OT.print_error(errmsg)
                break
            OT.print_warning(errmsg)

        if eof:
            done = True
        
        if skip: continue
        
        # save fasta
        fastaDB[label] = sequence

        # save processed species
        if species is None:
            species = getSpecies(label,specsep)
            if species not in speciesDB:
                speciesDB.append(species)
 
    return (errcode,stype,fastaDB,speciesDB)

def readFastaVect(inpf,specsep,labels=None,verbose=False):

    # loop over all fasta
    done=False
    errcode=0
    errmsg=None
    stype=None
    species=None
    fastaDB=[]
    speciesDB=[]

    while not done:
        # get next fasta
        (eof,label,descr,sequence)=getNext(inpf)
        # parse result
        (done,skip,errmsg,stype)=parseResult(stype,eof,label,descr,sequence,specsep,labels,verbose)
        if errmsg is not None:
            if not skip:
                errcode=1
                OT.print_error(errmsg)
                break
            OT.print_warning(errmsg)

        if eof:
            done = True
        
        if skip: continue

        fastaDB.append((label,sequence))

        # save processed species
        if species is None and specsep is not None:
            species = getSpecies(label,specsep)
            if species not in speciesDB:
                speciesDB.append(species)

    return (errcode,stype,fastaDB,speciesDB)

def getFastaLabels( fastaDB ):
    labels=[]
    for lbl,seq in fastaDB.items():
        labels.append(lbl)
    return labels


def makeType( typestr ):
    if typestr is None : return None

    # make upper case and select the first four chars
    rval=typestr.upper()
    rval=rval[:4]
    if not (rval == 'PROT' or rval == 'NUCL'):
        rval=None
    return rval
        

def getType( sequence, verbose ):
    typefs=None
    found=False
    sequence = sequence.upper()
    nnucl=0
    nprot=0
    ncommon=0
    nbad=0
    ntot=len(sequence)
    for c in sequence:
        if ( c in SeqLettersNucl ):
            nnucl=nnucl+1
            if nprot>0:
                if verbose:
                    OT.print_error(('Found an illegal character %c in a PROT sequence')%(c))
                nbad=nbad+1
                break
            continue
        if ( c in SeqLettersCommon ):
            ncommon=ncommon+1
            continue
        if ( c in SeqLettersMasked ):    continue
        if ( c in SeqLettersStopCodon ): continue

        if ( c in SeqLettersProt ) or ( c in SeqLettersProtAmbig ):
            nprot=nprot+1
            if nnucl>0:
                if verbose:
                    OT.print_error(('Found an illegal character %c in a NUCL sequence')%(c))
                nbad=nbad+1
                break
                
#            if found and not typefs=='PROT':
#                found=False
#                if verbose:
#                    OT.print_error(('Found an illegal character %c in a NUCL sequence')%(c))
#                break
#            typefs='PROT'
#            found=True
        else:
            nbad=nbad+1
#            typefs=None
#            found=False
#            if verbose:
#                OT.print_error(('Found an illegal character %c in a sequence')%(c))
#            break
        # stop search if found
#        if found: break


    # figure out what we have found
    # nbad>0 -> BAD
    #
    # ncommon nnucl nprot type
    #  >=0      0    >0    P
    #  > 0      0     0    N (ambiguous)
    #  >=0     >0     0    N
    #  >=0     >0    >0    Bad!

    if verbose:
        OT.print_info('Nbad, Nnucl, Nprot, Ncommon = %d, %d, %d, %d'%(nbad,nnucl,nprot,ncommon))

    # return if bad chrs are found
    if nbad>0:
        found=False
        typefs='BADT'
    elif nnucl>0 and nprot>0:
        found=False
        typefs='BADT'
    elif nnucl==0 and nprot>0:
        # make sure it's float otherwise it will fail in python 2
        fprot=float(nprot)/float(ntot)
        fprotCut=0.10
        # note the cut off in fraction is arbitrary...
        if fprot<fprotCut:
            # too low fraction of proteins - could be an error in the sequence
            if verbose:
                OT.print_error(('Ambiguous type due to low protein fraction (%6.3f < %6.3f)')%(fprot,fprotCut))
            found=False
            typefs='BADT'
        else:
            found=True
            typefs='PROT'
    elif nnucl>0 and nprot==0:
        found=True
        typefs='NUCL'
    elif ncommon>0 and nnucl==0 and nprot==0:
        # ambiguous but most likely NUCL
        found=True
        typefs='NUCL'

    if verbose:
        pt=typefs
        if pt is None: pt='NONE'
        OT.print_info('type is %s'%pt)

    return found,typefs

def checkSequence( sequence,verbose ):
    sequence = sequence.upper()
    nnucl=0
    ncommon=0
    nprot=0
    found=False
    errmsg=None
    for c in sequence:
        if ( c in SeqLettersNucl ):
            nnucl=nnucl+1
            continue
        if ( c in SeqLettersCommon ):
            ncommon=ncommon+1
            continue
        if ( c in SeqLettersMasked ):    continue
        if ( c in SeqLettersStopCodon ): continue

        if ( c in SeqLettersProt ) or ( c in SeqLettersProtAmbig ):
            nprot=nprot+1
        else:
            errmsg=('Found an illegal character %c in a sequence')%(c)
            if verbose:
                OT.print_error(errmsg)
            break

    return errmsg

def getTypeFile( fasta, verbose ):
    typefs=None
    found=False
    while not found:
        # get next fasta entry
        (eof,label,descr,sequence)=getNext(fasta)
        # if sequence is empty means that an illegal line is found - stop!
        if sequence is None:
            break
        # non empty sequence - check type
        if len(sequence)>0 and not found:
            found,typefs=getType(sequence, verbose)
            if typefs=='BADT':
                if verbose:
                    OT.print_error('Found bad type for label %s'%label)
                continue

    fasta.seek(0,0)
    return found,typefs

def getTypeFileOLD( fasta, verbose ):
    typefs='X'
    found=False
    for line in fasta:
        line = line.strip().rstrip('\n')
        if len(line)>0:
            if line[0] in ('#','>'): continue
            line = line.upper()
            found,typefs=getType(line, verbose)
            # stop if type is established
            if found:
                break
            # stop if bad line is encountered
            if typefs=='X':
                if verbose:
                    OT.print_error(('Unexpected character in string: %s')%(line))
                break
            if typefs=='UNDEF':
                if verbose:
                    OT.print_error(('Undefined: %s')%(line))
                break

    fasta.seek(0,0)
    return found,typefs

def maskStopCodons( sequence, lchr ):
    """replace stop codons with masked ('X' or 'N')"""
    rval=re.sub('[%s]'%(RegexCodons),lchr,sequence)
    return rval.rstrip(lchr)

def removeTrailingMasked( sequence, lchr ):
    return sequence.rstrip(lchr)

def cleanupFasta( fasta, target, fstype, verbose=False ):
    # make sure type is valid
    fstype=makeType(fstype)
    if fstype is None: return 2

    done=False
    stype=None
    nwarns=0
    nwmax=20
    nwxtras=0
    errcode=0
    errmsg=None
    # somewhat unclean...
    maskedchr=getMasked(fstype)
    if maskedchr is None: return 3

    while not done:
        (eof,label,descr,sequence)=getNext(fasta)

        (done,skip,errmsg,stype)=parseResult(stype,eof,label,descr,sequence,None,None,verbose)
        if errmsg is not None:
            if not skip:
                errcode=1
                OT.print_error(errmsg)
                break
            OT.print_warning(errmsg)

        if eof:
            done = True
        
        if skip: continue

        newseq=replaceLowerCaseMasked(sequence,maskedchr)
        newseq=maskStopCodons( newseq,maskedchr )

        if descr is None:
            descr=''
        if len(descr)>0: descr=' '+descr

        lblstr = ('>{0}{1}\n').format(label,descr)
        target.write(lblstr.encode('utf-8'))
        target.write(newseq.encode('utf-8'))
        target.write('\n'.encode('utf-8'))

        if not newseq==sequence:
            if nwarns==nwmax: OT.print_warning('above warning message supressed')
            if nwarns<nwmax:  OT.print_warning(('stop codons replaced with \'X\' in %s')%(label))
            else: nwxtras=nwxtras+1
            nwarns=nwarns+1

    if nwxtras>0: OT.print_warning(('supressed %d warnings (total = %d)')%(nwxtras,nwarns))
    return errcode


def statSequence( sequence ):
    hist={}
    for c in sequence:
        if c not in hist:
            hist[c] = 1
        else:
            hist[c] = hist[c]+1
    return (len(sequence),hist)

def statCharSequence( sequence, char ):
    nc=0
    for c in sequence:
        if c==char:
            nc = nc+1
    return nc

def checkMaskedSequence( sequence, typefs ):
    masked=getMasked(typefs)
    nmasked = statCharSequence(sequence,masked)
    fmasked = float(nmasked)/float(len(sequence))
    return fmasked

def longestSequence( sequence, masked, lencut=None ):
    """get longest consecutive unmasked sequence"""
    if masked is None: return len(sequence)
    nmax=0
    n=0
    for c in sequence:
        if c==masked:
            # if we have found lencut length, quit
            if lencut is not None:
                if nmax>=lencut: break
            n=0
        else:
            n=n+1
            if n>nmax: nmax=n

    return nmax

def checkMaskedFasta( fasta, typefs, minf, verbose ):
    data=[]
    done=False

    eof=False
    label=None
    description=None
    sequence = None

    while not done:
        (eof,label,description,sequence)=getNext(fasta)
        if eof:
            done=True
            continue
        elif label is None:
            if verbose:
                print ('could not find a valid fasta entry')
            errcode=3
            done=True
            continue

        if typefs is None:
            found,typefs = getType(sequence, verbose )

        if not found:
            OT.print_error(('unable to establish type of label <%s>') % (label))
            continue

        # check sequence and get fraction of masked bases
        fmasked=checkMaskedSequence( sequence, typefs )
        if fmasked<minf: continue

        if verbose: print(('%s    %6.3f')%(label,fmasked))
        data.append(fmasked)

    return data

def checkFastaType( fasta, expected, quick, specsep, verbose ):
    """
    Checks the fasta types and labels in a file
    error codes:
    0 - OK file
    1 - invalid label
    2 - bad fasta def found
    3 - inconsistent type
    4 - wrong type
    5 - empty file
    6 - type not established
    """
    typefsFound=None
    nok=0
    nfasta=0
    errcode=0
    errmsg=None
    label=None
    descr=None
    sequence=None
    typefs=None
    done=False
    while not done:
        (eof,label,descr,sequence)=getNext(fasta)
        if verbose:
            OT.print_info('checking label: %s'%label)

        (done,skip,errmsg,typefs)=parseResult(typefs,eof,label,descr,sequence,specsep,None,verbose)
        if errmsg is not None:
            if not skip:
                errcode=1
                OT.print_error(errmsg)
                break
            OT.print_warning(errmsg)

        if eof:
            done = True
        
        if skip: continue

        if verbose:
            OT.print_info('found type: %s'%typefs)

        nfasta=nfasta+1
            
        if typefs is None:
            # a bad fasta definition is found!
            errcode=2
            break
        elif typefsFound is None:
            # set found fasta type
            typefsFound=typefs
        elif not typefsFound==typefs:
            # inconsistent type
            errcode=3
            break

        # check if type is as expected
        if expected is not None and typefsFound is not None:
            if (not expected==typefsFound) and errcode==0:
                errcode=4

        if not quick and not eof:
            done=False

    # fasta file is empty
    if errcode==0 and nfasta==0:
        errcode=5

    # no proper type is found - should not happen (see errcode=2)
    if errcode==0 and typefsFound is None:
        errcode=6


    fasta.seek(0,0)

    return errcode,label,typefsFound


def readBlastclust( infile, representative, clsize ):
    done=False
    labels=[]
    myline=''
    while not done:
        myline = infile.readline()
        if len(myline)==0:
            done=True
            continue

        tokens=myline.split()
        # as soon as a cluster of a min size is found
        # skip the rest as they are sorted by cluster size
        if len(tokens)<clsize:
            done=True
            continue

        # save cluster, check if only representative (first in list) should be saved
        for t in tokens:
            labels.append(t)
            if representative: break

    return labels
