#!/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 searches for cd-hit cases where masked out regions have been aligned.
It was corrected in the latest version of cd-hit.
Pipeline usage: -
'''
from __future__ import print_function

#
# This searches for cd-hit cases where masked out regions have been aligned.
# It was corrected in the latest version of cd-hit.
#
import sys
import argparse
import fasta_tools as FT
import orthologer_tools as OT

badKeys={}

def cmp_seq(seqA,seqB,masked,debug=False):
    if seqA is None : return None,None
    if seqB is None : return None,None

    if debug:
        print('DEBUG: A')
        print(seqA)
        print('DEBUG: B')
        print(seqB)

    nmaskedA=FT.statCharSequence(seqA,masked)
    nmaskedB=FT.statCharSequence(seqB,masked)
    fmaskedA=float(nmaskedA)/float(len(seqA))
    fmaskedB=float(nmaskedB)/float(len(seqB))

    if debug:
        print((('nmasked A: %d, f=%8.4f')%(nmaskedA,fmaskedA)))
        print((('nmasked B: %d, f=%8.4f')%(nmaskedB,fmaskedB)))

    if not len(seqA)==len(seqB):
        # if sequence lengths are different, it is due to gaps
        # gap info is not given so estimate by count X's in each sequence
        # take the shortest masked region for our fX
        if nmaskedA<nmaskedB:
            fX = fmaskedA
        else:
            fX = fmaskedB
        # flag the presence of gaps/inserts by making the values negative
        fX = -fX
    else:
        fX = min(fmaskedA,fmaskedB)

    if debug:
        print((('f(masked) : %8.4f')%(fX)))

    return fX


def read_clstr(inpfile,stype,fxmax,fastaDB):
    """Read the cluster file output from cd-hit"""
    try:
        inpf = open(inpfile,"r")
    except IOError:
        print(( 'ERROR: file open failed for %s' % (inpfile) ))
        sys.exit(-1)

    #    print '# read clstr file: '+inpfile
    # cluster id
    clid=0
    # list of keys in cluster
    allkeys=[]
    # main key of cluster
    mainkey = None

    # character for masked base
    masked = FT.getMasked(stype)

    # loop over all lines in input file
    # expects a line like '1       59nt, >Hsap:71763444... at 1:58:28:85/+/98%'
    for line in inpf:
        # cleanup line
        line = line.strip().rstrip('\n')
        # skip empty or commented lines
        if (len(line)==0) or (line[0]=='#'): continue
        if line.startswith('>Cluster'):
            if not len(allkeys)==0:
            # process previous cluster
            # get sequence of main
                if mainkey in fastaDB:
                    mainseq = fastaDB[mainkey]
                    for keydata in allkeys:
                        key   = keydata[0]
                        ali   = keydata[1].split(':')
                        #pid = float(keydata[2])
                        if not key in fastaDB:
                            print(( ('WARNING: missing key in fasta file %s') % key ))
                            continue
                        keyseq = fastaDB[key]
                        mainstart= int(ali[2])
                        mainend  = int(ali[3])
                        keystart = int(ali[0])
                        keyend   = int(ali[1])

                        if mainend>len(mainseq):
                            print(( ('ERROR: %s main end pos beyond end of string %d > %d') % (key, mainend,len(mainseq) ) ))
                        if mainstart>len(mainseq):
                            print(( ('ERROR: %s main start pos beyond end of string %d > %d') % (key, mainstart,len(mainseq) ) ))

                        if keyend>len(keyseq):
                            print(( ('ERROR: %s key end pos beyond end of string %d > %d') % (key, keyend,len(keyseq) ) ))
                        if keystart>len(keyseq):
                            print(( ('ERROR: %s key start pos beyond end of string %d > %d') % (key, keystart,len(keyseq) ) ))

                        if mainstart>mainend:
                            mainseqA = mainseq[mainend-1:mainstart] #FT.complement(mainseq)
                        else:
                            mainseqA = mainseq[mainstart-1:mainend]

                        if keystart>keyend:
                            keyseqA = keyseq[keyend-1:keystart] #FT.complement(keyseq)
                        else:
                            keyseqA = keyseq[keystart-1:keyend]

                        # compare sequences
                        #   f  = N(non masked)/Length
                        #   fX = N(masked)/length
                        # if sizes are different:
                        #   f  = (100-fX)
                        #   fX = N(masked, shortest)/Length
                        #   both fX and f are negated to flag this case
                        debug=False
                        dbgkey=''
#                        dbgkey='320372:00060a'
#                        dbgkey='MDEST:bg7chp'
                        if key==dbgkey: debug=True
                        fX = cmp_seq(mainseqA,keyseqA,masked,debug)
                        debug=False

                        # flag presence of gaps
                        #hasgaps=False
                        f=0.0
                        if fX<0:
                            f=-1.0-fX
                            #hasgaps=True
                        else:
                            f=1.0-fX

                        if abs(fX)>fxmax:
                            global badkeys
                            # large fraction of X's -> save it as a bad one
                            #gapstr=''
                            #if hasgaps: gapstr='; GAPS'
                            # print('%s vs %s ; %5.1f  %5.1f  %5.1f %s' % (mainkey,key,f,fX,pid,gapstr) )
                            if mainkey not in badKeys:
                                badKeys[mainkey]=f
                            if key not in badKeys:
                                badKeys[key]=f

                            #print('%s ; %8.4f' % (mainkey,fX) )

            #clid = line.split()[1]
            mainkey = None
            allkeys = []
            #alilen = 0

        elif line.endswith('*'):
        # found a main key
            ldata = line.split()
            key = ldata[2].strip('>.');
            if not mainkey==None:
                print(( 'WARNING: got a second main key - first = %s, second = %s' % (mainkey,key) ))
            mainkey = key
        # ldata[1] = '45nt,'
            #alilen = int(ldata[1][0:-3])
        else:
        # split line according to whitespace
            ldata = line.split()
        # get key in list, remove leading '>' and trailing '...'
            key = ldata[2].strip('>.');
        # get alignment string e.g
        #   NUCL: '1:40:2:45/+/98%'
        #   PROT: '1:40:2:45/98%'
            alidata = ldata[4].split('/')
            # allkeys contains
            #   (key,alignment,pid)
            allkeys.append([key,alidata[0],alidata[-1].rstrip('%')])


def get_fasta_file(datadir,s):
    f=datadir+s+'.fs_masked'
    return f

def get_clstr_file(datadir,s):
    f=datadir+s+'.fs97.clstr'
    return f

def is_key_ok(key,fxmax,datadir,fastaDB,speciesDB):
    global badKeys
    if key in badKeys: return False,badKeys[key],fastaDB,speciesDB
    if key in fastaDB: return True,None,fastaDB,speciesDB
    # key is not in fastaDB or bad keys
    # check species
    species = FT.getSpecies(key)
    if species in speciesDB:
        # species has been read but not the given key!
        print('WARNING: fasta file for given species has been read but key is not found '+key)
        return True,None,fastaDB,speciesDB

    # read fasta file
    inpfile = get_fasta_file(datadir,species)
    try:
        ffasta=open(inpfile)
    except IOError:
        OT.print_error( ('failed opening fasta file %s')%(inpfile) )
        sys.exit(-1)

    (ec,stype,fastaDB,speciesDB) = FT.readFasta(ffasta,None)
    if ec>0:
        print(( 'ERROR: file open failed with ec = %d for species %s' % (ec,species) ))
        sys.exit(-1)

    # read cd-hit cluster file
    read_clstr(get_clstr_file(datadir,species),stype,fxmax,fastaDB)

    # recheck key
    if key in badKeys: return False,badKeys[key],fastaDB,speciesDB

    # all is well
    return True,None,fastaDB,speciesDB

def read_clusters(inpfile,datadir,fxmax):
    # read cluster file
    try:
        inpf = open(inpfile,"r")
    except IOError:
        print(( 'ERROR: file open failed for %s' % (inpfile) ))
        sys.exit(-1)

    clid=None
    keytype9=[]
    #spec2clid={}
    ntot=0
    fastaDB={}
    speciesDB=[]
    print(('#%-11s %4s %8s %8s %8s %20s %s') % ("clid","N(type9)","f(bad in cluster)", "sum(f)","sumf with gaps","main key","type 9 keys"))
    for line in inpf:
        # cleanup line
        line = line.strip().rstrip('\n')
        # skip empty or commented lines
        if (len(line)==0) or (line[0]=='#'): continue
        # clid  key     type    length  start   end     rawscore        normscore       evalue
        # get cluster data
        elems=line.split()
        clidnew=int(elems[0])
        key=elems[1]
        ctype=int(elems[2])

        if (not clidnew==clid) and (clid is not None):
            # a new cluster found - check previous cluster
            if len(keytype9)>0:
                sumf=0.0;
                sumnf=0.0;
                #ftype9 = float(ntot)/float(len(keytype9))
                worstkey=None
                worstnkey=None
                worstf=1000.0
                worstnf=-1000.0
                nbadn=0
                nbadp=0
                nbad=0
                badsp={}
                # loop over all type 9 keys
                for k9 in keytype9:

                    # key is OK if the fraction of X is low enough
                    keyok,f,fastaDB,speciesDB=is_key_ok(k9,fxmax,datadir,fastaDB,speciesDB)
                    if f is None: f=0.0
#                    print ('key=%s got f=%8.4f')%(k9,f)

                    if not keyok:
                        # high fraction of X
                        print(('%-16s %8.4f %8d') % (k9,f,clidnew))
                        species=FT.getSpecies(k9)
                        nbad=nbad+1
                        if species not in badsp: badsp[species]=0
                        badsp[species] = badsp[species]+1

#---------------------- if per species to clid map is requested
#                        if species not in spec2clid: spec2clid[species]={}
#                        if clid not in spec2clid[species]: spec2clid[species][clid]=0

#                        fadd=1.0-abs(f/100.0)
#                        spec2clid[species][clid]=spec2clid[species][clid]+fadd
#-------------------------------------------------------------------------------
                        if f<0:
                            nbadn=nbadn+1
                            sumnf = sumnf+(f/100.0)
                            if f>worstnf:
                                worstnf=f
                                worstnkey=k9
                        else:
                            nbadp = nbadp+1
                            sumf  = sumf+(f/100.0)
                            if f<worstf:
                                worstf=f
                                worstkey=k9
                            
                if nbad>0:
                    selkey=worstkey
                    if worstkey is None: selkey=worstnkey
                    badspstr=None
                    badspitems=list(badsp.items())
                    badspsrt=sorted(badspitems,key=lambda badspitems: -badspitems[1])
                    
                    for tp in badspsrt:
                        if badspstr is None: badspstr=''
                        else: badspstr=badspstr+';'
                        badspstr=badspstr+('%s,%d')%(tp[0],tp[1])
                    fbad=float(nbad)/float(len(keytype9))
                    print(('SUMMARY: %-12d %4d %8.2f %8.2f %8.2f %20s %s') % (clid,len(keytype9),fbad,sumf,sumnf,selkey,badspstr))


            ntot=0
            keytype9=[]

        # update data
        clid=clidnew
        ntot=ntot+1
        if ctype==9:
            keytype9.append(key)

#---------------------- if per species to clid map is requested
#    for sp2c in spec2clid:
#
#        sp2cstr=('%s ')%(sp2c)
#        sp2citems=spec2clid[sp2c].items()
#        sp2csorted=sorted(sp2citems,key=lambda sp2citems: -sp2citems[1])
#        for cid in sp2csorted:
#            sp2cstr=sp2cstr+('( %d %6.2f ) ')%(cid[0],cid[1]) #spec2clid[sp2c][cid])
#        print (('%s')%(sp2cstr))
#-------------------------------------------------------------------------------            
                
#-----------------------------------------------------------------
# MAIN

parser = argparse.ArgumentParser(description='Check for bad cd-hit output')

parser.add_argument('-c', '--cluster', dest='cluster', help='cluster file', metavar='FILE', required=True)
parser.add_argument('-d', '--datadir', dest='datadir', help='data directory with sequences', metavar='FILE', required=True)
parser.add_argument('--maxf', dest='maxf', help='maximum fraction masked bases allowed', type=float, default=0.6)

options = parser.parse_args()

# open fasta file and read data
read_clusters(options.cluster,options.datadir,options.maxf)

            
