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

'''
Takes the SELECT cluster output (.clstr) and scans for clusters with atleast two members.
These events are then output in a format readable by the merge inparalog code.
Pipeline usage: MAKEINPARSEL
'''
from __future__ import print_function

import sys
import orthologer_tools as OT

def coverageMin( seq1,seq2 ):
    len1 = seq1[1]-seq1[0]
    len2 = seq2[1]-seq2[0]
    maxstart = max(seq1[0],seq2[0])
    minend = min(seq1[1],seq2[1])
    overlap = minend-maxstart
    if overlap<=0:
        return 0;
    minlen = min(len1,len2)
    if minlen<1:
        return 0
    return float(overlap)/float(minlen)

def match( region1, region2, sequence, mincov ):
    rval = -1
    f1   = -1

    cov1 = coverageMin(region1,sequence)
    cov2 = coverageMin(region2,sequence)
    if cov1<mincov:
        cov1=0
    if cov2<mincov:
        cov2=0
    cmin = min(cov1,cov2)
    cmax = max(cov1,cov2)
    score=100.0
    if cmin>0:
        score = 100.0 - (100.0*cmin/cmax)
        f1 = cov1/(cov1+cov2)
        if score<5:
            # matches both regions - split
            rval = 3
        else:
            # unclear - majority vote
            if cov1>cov2:
                rval=1
            else:
                rval=2
    else:
        # clear classification
        if cov1>cov2:
            f1=1.0
            rval=1
        else:
            f1=0.0
            rval=2

    # note f2 == 1-f1
    return rval,f1


def loadSplit( splitfile ):
    data={}
    metadata={}
    nextMD = False
    nextData = 0
    Ndata = 0
    with open(splitfile,'r') as fh:
        for line in fh:
            line = line.strip().rstrip('\n')
            if len(line)==0:
                continue
            
            if line=="#METADATA":
                nextMD = True
                continue
            if line=="#DATA":
                nextData = 1;
                continue
            if nextMD:
                elems = line.split(' ')
                metadata['mincov'] = float(elems[0])
                metadata['minoverlap'] = float(elems[1])
                metadata['cutoff'] = float(elems[2])
                metadata['fkeep'] = float(elems[3])
                metadata['fgood'] = float(elems[4])
                metadata['mincount'] = float(elems[5])
                metadata['flimcut'] = float(elems[6])
                nextMD = False
                
            if nextData==1:
                Ndata = int(line)
                nextData = 2
                continue
            if nextData>1:
                # data line format
                # <key> <nregions> <count1> <start1> <end1> <count2> <start2> <end2> <minstart> <maxend> <maxcount> <totcount> 
                #  0     1          2        3        4      5        6        7      8          9        10         11
                nextData = nextData+1
                elems = line.split(' ')
                # allow only 2 regions
                nregs = int(elems[1])
                if nregs!=2:
                    continue
                key = elems[0]
                keydata = {}
                keydata['count1'] = int(elems[2])
                keydata['reg1'] = [ int(elems[3]), int(elems[4]) ]
                keydata['count2'] = int(elems[5])
                keydata['reg2'] = [ int(elems[6]), int(elems[7]) ]
                #keydata['minstart'] = int(elems[8])
                #keydata['maxend'] = int(elems[9])
                data[key] = keydata
    return Ndata,data,metadata

def overlap(seq,reg, refpos=None):
    if refpos is None:
        return [ max(seq[0],reg[0]), min(seq[1],reg[1]) ]
    regr = [ reg[0]-refpos, reg[1]-refpos ]
    return [ max(seq[0],regr[0]), min(seq[1],regr[1]) ]

def addkeysSplit(allkeys,mainkey,keysplitregion,Nsplits):
    mainsplit = None
    reg1 = None
    reg2 = None

    if mainkey in keysplitregion:
        mainsplit = keysplitregion[mainkey]
        reg1 = mainsplit['reg1']
        reg2 = mainsplit['reg2']

    outstr1 = None
    outstr2 = None
    for keydata in allkeys:
        key   = keydata[0]
        ali   = keydata[1].split(':')
        pid   = keydata[2]
        # sequences of this key and the main
        # sequences are starting at 1 in cdhit output
        # split file starts at 0 ...
        seqinp = [ int(ali[0])-1, int(ali[1])-1 ]
        seqkey = [ int(ali[2])-1, int(ali[3])-1 ]
        mainkey1 = mainkey
        mainkey2 = None
        key1 = key
        key2 = None
        seqkey1 = seqkey
        seqkey2 = None
        seqinp1 = seqinp
        seqinp2 = None
        if mainsplit is None:
            # no split
            Nsplits[1] = Nsplits[1]+1
        else:
            ec,f1 = match( reg1, reg2, seqkey,0.1 )
            Nsplits[ec+1] = Nsplits[ec+1]+1
            if ec==-1:
                continue
            if ec==1:
                # first region only
                mainkey2 = None
                key2 = None
                seqkey2 = None
                seqinp2 = None
                mainkey1 = '%s_00'%(mainkey)
                key1 = '%s_00'%(key)
                seqkey1 = seqkey
                seqinp1 = seqinp
            elif ec==2:
                # second region only
                mainkey1 = None
                key1 = None
                seqkey1 = None
                seqinp1 = None
                mainkey2 = '%s_01'%(mainkey)
                key2 = '%s_01'%(key)
                seqkey2 = seqkey
                seqinp2 = seqinp
            elif ec==3:
                # split!
                mainkey1 = '%s_00'%(mainkey)
                key1 = '%s_00'%(key)
                mainkey2 = '%s_01'%(mainkey)
                key2 = '%s_01'%(key)
                seqkey1 = overlap(seqkey,reg1)
                seqinp1 = overlap(seqinp,reg1,seqkey[0])
                seqkey2 = overlap(seqkey,reg2)
                seqinp2 = overlap(seqinp,reg2,seqkey[0])

        if mainkey1 is not None:
            if outstr1 is None:
                outstr1 = '{} - 0.0 '.format(mainkey1)
                #outstr1 = '%s:{'%(str(mainkey1))
            else:
                outstr1 += ' '
                #outstr2 += ','
            outstr1 += '{};{};{}:{}'.format(key1,pid,seqinp1[0]+1,seqinp1[1]+1)
            #outstr1 += '%s(%s;%d;%d)' % (str(key1),pid,seqinp1[0]+1,seqinp1[1]+1)
        if mainkey2 is not None:
            if outstr2 is None:
                outstr2 = '{} - 0.0 '.format(mainkey2)
                #outstr2 = '%s:{'%(str(mainkey2))
            else:
                outstr2 += ' '
                #outstr2 += ','
            outstr2 += '{};{};{}:{}'.format(key2,pid,seqinp2[0]+1,seqinp2[1]+1)
            #outstr2 += '%s(%s;%d;%d)' % (str(key2),pid,seqinp2[0]+1,seqinp2[1]+1)

#        if mainkey1 is not None:
#            if outstr1 is None:
#                outstr1 = '%s:{'%(str(mainkey1))
#            else:
#                outstr1 += ','
#            outstr1 += '%s(%s;%d;%d)' % (str(key1),pid,seqinp1[0]+1,seqinp1[1]+1)
#        if mainkey2 is not None:
#            if outstr2 is None:
#                outstr2 = '%s:{'%(str(mainkey2))
#            else:
#                outstr2 += ','
#            outstr2 += '%s(%s;%d;%d)' % (str(key2),pid,seqinp2[0]+1,seqinp2[1]+1)    
#    if outstr1 is not None:
#        outstr1 += '}'
#    if outstr2 is not None:
#        outstr2 += '}'
    return outstr1,outstr2,Nsplits



def addkeys(allkeys,mainkey,mlen):
    outstr = None
    try:
        mainlen = int(mlen.split('aa,')[0])
    except ValueError:
        mainlen = 0

    for keydata in allkeys:
        key    = keydata[0]   # gene id
        slen   = keydata[1]   # eg: 1101aa,
        ali    = keydata[2].split(':') # alignment info
        pid    = keydata[3]   # percentage id
        try:
            seqlen = int(slen.split('aa,')[0])
        except ValueError:
            seqlen = mainlen
        
        # sequences of this key and the main
        # sequences are starting at 1 in cdhit output
        # split file starts at 0 ...
        seqinp = [ int(ali[0])-1, int(ali[1])-1 ]
        seqkey = [ int(ali[2])-1, int(ali[3])-1 ]

        # if sequence is
        # - [0,0] -> [0,len-1]
        # - [x,0] -> [x,len-1]
        if not (seqinp[0]<seqinp[1]):
            seqinp[1]=seqlen-1

        if not (seqkey[0]<seqkey[1]):
            seqkey[1]=mainlen-1
        
        if mainkey is not None:
            if outstr is None:
                outstr = '{} - 0.0 '.format(mainkey)
            else:
                outstr += ' '
            outstr += '{};{};{};{};{};{}'.format(key,pid,seqkey[0]+1,seqkey[1]+1,seqinp[0]+1,seqinp[1]+1)
            
    return outstr

#-----------------------------------------------------------------
# MAIN
# check args
if len(sys.argv)<2:
    print("usage: "+sys.argv[0]+"<clstr file> [<split file>]")
    sys.exit(1)

# set args
splitfile = ""
inpfile = sys.argv[1]
if len(sys.argv)>2:
    splitfile = sys.argv[2]
    
# set split data
nsplits=0
splitdata=None
splitmd={}

if len(splitfile)>0:
    nsplits,splitdata,splitmd = loadSplit( splitfile )
    print('# splits read {} keys from {}'.format(nsplits,splitfile))
    if nsplits==0:
        splitdata = None

# open cd hit cluster file and read data
try:
    inpf = open(inpfile,"r")
except IOError:
    OT.print_error('file open failed for %s' % (inpfile))
    sys.exit(-1)

# list of keys in cluster
allkeys=[]

# main key of cluster - set whenever a line that ends with * is found
mainkey = None
mainlen = None

Nsplits=[0,0,0,0,0]

# loop over all lines in input file
# expects a line like '1       59nt, >Hsap:71763444... at 1:58:28:85/+/98%'
# the sequence position fmt is <start:end of this>:<start:end of *>
for line in inpf:
    line = line.strip().rstrip('\n')
    if len(line)>0 and (not line[0]=='#'):
        if line.startswith('>Cluster'):
            if len(allkeys)>0:
                # save previous cluster
                if splitdata is None:
                    outstr1 = addkeys( allkeys, mainkey, mainlen )
                    outstr2 = None
                else:
                    outstr1,outstr2,Nsplits = addkeysSplit( allkeys, mainkey, splitdata, Nsplits )

                if outstr1 is not None:
                    print('%s' % (outstr1))
                if outstr2 is not None:
                    print('%s' % (outstr2))
            
            mainkey = None
            mainlen = None
            allkeys = []

        elif line.endswith('*'):
            # found a main key
            ldata = line.split()
            # get key
            # two different versions exists of cdhit:
            # with or without space between '>' and label
            key = ldata[2].strip('>.');
            if len(key)==0:  key = ldata[3].strip('>.');
            if not mainkey==None:
                print("WARNING: got a second main key - first = %s, second = %s" % (mainkey,key))
            mainkey = key
            mainlen = ldata[1]
        else:
            # first key of new cluster
            # fmt:
            # 0       2278aa, > AFLOR:001f10... at 1:2278:23:2317/98.64%
            #
            # split line according to whitespace
            ldata = line.split()
            # get key in list, remove leading '>' and trailing '...'
            # again, two versions of format (see above)
            key = ldata[2].strip('>.');
            if len(key)==0:  key = ldata[3].strip('>.');

            # get alignment string e.g
            #   NUCL: '1:40:2:45/+/98%'
            #   PROT: '1:40:2:45/98%'
            lind=4
            if ldata[lind]=="at": lind=5
            alidata = ldata[lind].split('/')
            allkeys.append([key,ldata[1],alidata[0],alidata[-1].rstrip('%')])

# print any remaining keys
if len(allkeys)>0:
    
    # save previous cluster
    if splitdata is None:
        outstr1 = addkeys( allkeys, mainkey, mainlen )
        outstr2 = None
    else:
        outstr1,outstr2,Nsplits = addkeysSplit( allkeys, mainkey, splitdata, Nsplits )
    
    if outstr1 is not None:
        print('%s' % (outstr1))
    if outstr2 is not None:
        print('%s' % (outstr2))

if splitdata is not None:
    stxt = [ 'bad','none','reg1','reg2','split' ]
    for i in range(0,len(Nsplits)):
        print('# splits %-6s : %d'%(stxt[i],Nsplits[i]))
    
