#!/usr/bin/env python
# -*- coding: utf-8 -*-

from matching import matchAtoms
from atom import atomIsDonor, atomIsFromLigand, atomIsFromBackbone, getOneLetterResudueCode, atomIsFromProtein, atomIsFromDNAorRNA
import copy
import pickle

def getPenalties(atomsList):
    
    P=0
    N=len(atomsList)
        
    atom1_1 = atomsList[0][0]
    atom2_1 = atomsList[0][1]
            
    id1 = atom1_1.resid
    id2 = atom2_1.resid
    
    n = 1
        
    for i in xrange(1, len(atomsList)):
        
        atom1_i  = atomsList[i][0]
        atom2_i  = atomsList[i][1]
            
        if ((atom1_i.resid != id1 + 1) and (atom1_i.resid != id1 - 1)) or ((atom2_i.resid != id2 + 1) and (atom2_i.resid != id2 - 1)): 
            P = P + 1./n
            n = 1
            
        else:
            n += 1
        
        id1 = atom1_i.resid
        id2 = atom2_i.resid
    
    P = P + 1./n
    
    
    return P, N

def penalizeSequence(pairsOfAtoms):
    
    reducedPairsOfAtoms = eliminateSideChains(pairsOfAtoms)
    atomsList = quicksortPairsByChainThenID(reducedPairsOfAtoms, 0)
    
    atomsListProtein = []
    atomsListDNAorRNA = [] 
    for atomPair in atomsList:
        if atomIsFromProtein(atomPair[0]):
            atomsListProtein.append(atomPair)
        elif atomIsFromDNAorRNA(atomPair[0]):    
            atomsListDNAorRNA.append(atomPair)
    
    P = 0
    N = 0
    if atomsListProtein:
        P_prot, N_prot = getPenalties(atomsListProtein)
        P += P_prot
        N += N_prot
    if atomsListDNAorRNA:
        P_nucl, N_nucl = getPenalties(atomsListDNAorRNA)
        P += P_nucl
        N += N_nucl
    
    return P, N

def getResidueIdsOfAtomList(itemList):
    
    sorted_list = quicksortListByChainThenID(itemList)

    residues = []
    for atom in sorted_list:
        if atomIsFromProtein(atom) or atomIsFromDNAorRNA(atom):
            residues.append(atom.resid)
            
    return residues

def getSequenceOfAtomList(itemList):
    
    sorted_list = quicksortListByChainThenID(itemList)
    
    sequence = []
    for atom in sorted_list:
        if atomIsFromProtein(atom) or atomIsFromDNAorRNA(atom):
            if getOneLetterResudueCode(atom) != None:
                aa = getOneLetterResudueCode(atom)
                sequence.append(aa)
            else:
                sequence.append('X')
    
    return sequence
    
def getSequenceOfAlignment(atomsList):
    
    atomsList_sorted = quicksortPairsByChainThenID(atomsList, 0)
    
    atom1_1 = atomsList[0][0]
    id1 = atom1_1.resid
    aa1 = getOneLetterResudueCode(atom1_1)
    
    atom2_1 = atomsList[0][1]
    id2 = atom2_1.resid
    aa2 = getOneLetterResudueCode(atom2_1)
    
    aaList1 = []
    resList1 = []
    
    aaList2 = []
    resList2 = []
       
    for i in xrange(1, len(atomsList)):
        
        atom1_i  = atomsList[i][0]
        atom2_i  = atomsList[i][1]
        
        if (getOneLetterResudueCode(atom1_i) != None and getOneLetterResudueCode(atom2_i) != None):
            
            aaList1.append(getOneLetterResudueCode(atom1_i))
            resList1.append(atom1_i.resid)
            
            aaList2.append(getOneLetterResudueCode(atom2_i))
            resList2.append(atom2_i.resid)
            
    return aaList1, resList1, aaList2, resList2

def writeSequenceToFile(site1, site2, outputPath, maxDist):

    pairsOfAtoms = matchAtoms(site1, site2, maxDist)
    reducedPairsOfAtoms = eliminateSideChains(pairsOfAtoms)   
    
    atomsList = quicksortPairsByChainThenID(reducedPairsOfAtoms, 0)
    
    atomsListProtein = []
    atomsListDNAorRNA = [] 
    for atomPair in atomsList:
        if atomIsFromProtein(atomPair[0]):
            atomsListProtein.append(atomPair)
        elif atomIsFromDNAorRNA(atomPair[0]):    
            atomsListDNAorRNA.append(atomPair)
    
    sequence_list = []
    sequence_list_prot = []
    sequence_list_nucl = []
    
    outFile = open(outputPath, 'w')
    #    
    querySiteName = '{0:<7}{1:<4}'.format('Query:', site1.name)
    sbjctSiteName = '{0:<7}{1:<4}'.format('Sbjct:', site2.name)    
    #
    outFile.writelines(querySiteName)
    outFile.writelines('\n')
    outFile.writelines(sbjctSiteName)
    outFile.writelines('\n')
    
    if atomsListProtein:
        outFile.writelines('\n')
        outFile.writelines('Protein alignment:\n')
        
        sequence_list = writeLinesOfSequenceAlignment(atomsListProtein, outFile, sequence_list)
        sequence_list_prot = copy.deepcopy(sequence_list)
        sequence_list = []
        
    if atomsListDNAorRNA:
        outFile.writelines('\n')
        outFile.writelines('DNA/RNA alignment:\n')
   
        sequence_list = writeLinesOfSequenceAlignment(atomsListDNAorRNA, outFile, sequence_list)
        sequence_list_nucl = copy.deepcopy(sequence_list)
    outFile.close()
    
    return sequence_list_prot, sequence_list_nucl

def getLinesOfSequenceAlignment(atomsList):
    
    #
    atom1_1 = atomsList[0][0]
    chain1 = atom1_1.chain
    id1 = atom1_1.resid
    aa1 = getOneLetterResudueCode(atom1_1)
    #
    atom2_1 = atomsList[0][1]
    chain2 = atom2_1.chain
    id2 = atom2_1.resid   
    aa2 = getOneLetterResudueCode(atom2_1)
    #
    
    #
    chainList1 = []
    chainList1.append('{0:<4}'.format(chain1))
    idList1 = []
    idList1.append('{0:<4}'.format(id1))
    aaList1 = []
    if atomIsFromLigand(atom1_1):    
        aaList1.append('{0:<1}{1:<3}'.format(aa1, '*'))
    else:   
        aaList1.append('{0:<4}'.format(aa1))
    #
    pipeList = []
    #
    chainList2 = []
    chainList2.append('{0:<4}'.format(chain2))
    idList2 = []
    idList2.append('{0:<4}'.format(id2))
    aaList2 = []
    if atomIsFromLigand(atom2_1):    
        aaList2.append('{0:<1}{1:<3}'.format(aa2, '*'))
    else:   
        aaList2.append('{0:<4}'.format(aa2))
    #
    
    if aaList1[0] == aaList2[0]:
        pipeList.append('|  ')
    else:
        pipeList.append('   ')
    
    k = 0    
    for i in xrange(1, len(atomsList)):
        
        atom1_i  = atomsList[i][0]
        atom2_i  = atomsList[i][1]

        if (getOneLetterResudueCode(atom1_i) != None and getOneLetterResudueCode(atom2_i) != None):
            flag = False
            k=k+1
            # The first atom from a pair         
            if atom1_i.chain != chain1:
                chainList1.append(' | ')
                idList1.append(' | ')
                aaList1.append(' | ')
                #
                pipeList.append('   ')
                #
                chainList2.append(' | ')
                idList2.append(' | ')
                aaList2.append(' | ')                
                chainList1.append('{0:<4}'.format(atom1_i.chain))
                chain1 = atom1_i.chain
                flag = True
                k=k+1           
            else:
                chainList1.append(' '*4)
                
            if (atom1_i.resid != id1 + 1) and (atom1_i.resid != id1 - 1):
                if flag == False:                
                    temp = chainList1[k]
                    chainList1[k] = '   '
                    chainList1.append(temp)
                    idList1.append(' | ')
                    aaList1.append(' | ')
                    #
                    pipeList.append('   ')
                    #
                    chainList2.append('   ')
                    idList2.append(' | ')
                    aaList2.append(' | ')
                    k=k+1
                    flag=True
                
                idList1[k-2] = '{0:<4}'.format(atomsList[i-1][0].resid)
                idList1.append('{0:<4}'.format(atom1_i.resid))
                id1 = atom1_i.resid
            else:
                idList1.append('{0:<4}'.format(atom1_i.resid))
                id1 = atom1_i.resid
            
            if atomIsFromLigand(atom1_i):   
                aaList1.append('{0:<1}{1:<3}'.format(getOneLetterResudueCode(atom1_i), '*'))
            else:
                aaList1.append('{0:<4}'.format(getOneLetterResudueCode(atom1_i)))
            
            #The second atom from a pair        
            if atom2_i.chain != chain2:
                #If chain is changed in second sequence, pipe is inserted in both sequenses and first sequence is shifted
                if flag == False:                
                    temp = chainList1[k]
                    chainList1[k] = ' | '
                    chainList1.append(temp)
                    temp = idList1[k]
                    idList1[k] = ' | '    
                    idList1.append(temp)                    
                    temp = aaList1[k]
                    aaList1[k] = ' | '
                    aaList1.append(temp) 
                    chainList2.append(' | ')
                    idList2.append(' | ')
                    aaList2.append(' | ')
                    pipeList.append('   ')
                    k=k+1
                    flag = True
                else:
                    chainList1[k-1] = ' | '
                    chainList2[k-1] = ' | '
                    
                chainList2.append('{0:<4}'.format(atom2_i.chain))
                chain2 = atom2_i.chain
            else:
                chainList2.append(' '*4)
            
            aaTemp = aaList1[k]    
            if (atom2_i.resid != id2 + 1) and (atom2_i.resid != id2 - 1):
                if flag == False:                
                    temp = chainList1[k]
                    chainList1[k] = '   '
                    chainList1.append(temp)
                    temp = idList1[k]
                    idList1[k] = ' | '
                    idList1.append(temp)                    
                    temp = aaList1[k]
                    aaList1[k] = ' | '
                    aaList1.append(temp) 
                    temp = chainList2[k]
                    chainList2[k] = '   '
                    chainList2.append(temp)
                    idList2.append(' | ')
                    aaList2.append(' | ')
                    pipeList.append('   ')
                    k=k+1
                    flag = True
                    
                idList2[k-2] = '{0:<4}'.format(atomsList[i-1][1].resid)
                idList2.append('{0:<4}'.format(atom2_i.resid))
                id2 = atom2_i.resid
            else:
                idList2.append('{0:<4}'.format(atom2_i.resid))
                id2 = atom2_i.resid
            
            if atomIsFromLigand(atom2_i):    
                aaList2.append('{0:<1}{1:<3}'.format(getOneLetterResudueCode(atom2_i), '*'))
            else:
                aaList2.append('{0:<4}'.format(getOneLetterResudueCode(atom2_i)))
            
            if aaTemp == aaList2[-1]:
                pipeList.append(' |  ')
            else:
                pipeList.append('    ')
    
    return chainList1, idList1, aaList1, chainList2, idList2, aaList2, pipeList


def writeLinesOfSequenceAlignment(atomsList, outFile, sequence_list):

    chainList1, idList1, aaList1, chainList2, idList2, aaList2, pipeList = getLinesOfSequenceAlignment(atomsList)
    #
    
    sequence_list.append((chainList1, idList1, aaList1, pipeList, aaList2, idList2, chainList2))
    
    outFile.writelines('\n')   
    
    chainList1.insert(0, 'Chain:  | ')
    outFile.writelines(chainList1)
    outFile.writelines('\n')
    idList1.insert(0, 'ResID:  | ')
    outFile.writelines(idList1)
    outFile.writelines('\n')
    aaList1.insert(0, 'Query:  | ')
    outFile.writelines(aaList1)
    outFile.writelines('\n')
    
    pipeList.insert(0, '          ')
    outFile.writelines(pipeList)
    outFile.writelines('\n')
    
    aaList2.insert(0, 'Target: | ')
    outFile.writelines(aaList2)
    outFile.writelines('\n')
    idList2.insert(0, 'ResID:  | ')
    outFile.writelines(idList2)
    outFile.writelines('\n')
    chainList2.insert(0, 'Chain:  | ')
    outFile.writelines(chainList2)
    outFile.writelines('\n')
    
    return sequence_list

def eliminateSideChains(listOfAtomPairs):
    'Takes a list of all pairs and remains carbon alpha of protein residues and C1 of DNA/RNA residues'
    
    reducedListOfAtomPairs = [pairOfAtoms for pairOfAtoms in listOfAtomPairs if atomIsFromBackbone(pairOfAtoms[0])]
    
    return reducedListOfAtomPairs

#LIST OF PAIRS SORTING

def quicksortPairsByChainThenID(pairList, pairID):
    
    chainIDsortedPairs = quicksortPairsByChain(pairList, pairID)
    
    blocksOfChains = []
    chainBlock = []
    
    prev_item = chainIDsortedPairs[0]
    for next_item in chainIDsortedPairs:
        if prev_item[pairID].chain == next_item[pairID].chain:
            chainBlock.append(next_item)
            prev_item = next_item
        else:
            sequence = quicksortPairsByResID(chainBlock, pairID)
            blocksOfChains.append(sequence)
            chainBlock = []
            chainBlock.append(next_item)
            prev_item = next_item
            
    sequence = quicksortPairsByResID(chainBlock, pairID)
    blocksOfChains.append(sequence)

    pairsSortedByChainThenID = []
    for block in blocksOfChains:
        pairsSortedByChainThenID.extend(block)
    
    return pairsSortedByChainThenID    


def quicksortPairsByResID(pairList, pairID):
    'Takes a list of atom pairs and sorts all atoms by residue ID of atoms with a given ordinal number'
    
    if len(pairList) <= 1: return pairList
    
    pivot = pairList.pop()
    before = [pair for pair in pairList if pair[pairID].resid <= pivot[pairID].resid]
    after = [pair for pair in pairList if pair[pairID].resid > pivot[pairID].resid]
    
    
    return quicksortPairsByResID(before, pairID) + [pivot] + quicksortPairsByResID(after, pairID)


def quicksortPairsByChain(pairList, pairID):

    if len(pairList) <= 1: return pairList
    
    pivot = pairList.pop()
    before = [pair for pair in pairList if pair[pairID].chain <= pivot[pairID].chain]
    after = [pair for pair in pairList if pair[pairID].chain > pivot[pairID].chain]
    
    
    return quicksortPairsByChain(before, pairID) + [pivot] + quicksortPairsByChain(after, pairID)

#LIST SORTING

def quicksortListByChainThenID(itemList):

    chainIDsortedList = quicksortListByChain(itemList)
    
    blocksOfChains = []
    chainBlock = []
    
    prev_item = chainIDsortedList[0]
    for next_item in chainIDsortedList:
        if prev_item.chain == next_item.chain:
            chainBlock.append(next_item)
            prev_item = next_item
        else:
            sequence = quicksortListByResID(chainBlock)
            blocksOfChains.append(sequence)
            chainBlock = []
            chainBlock.append(next_item)
            prev_item = next_item
            
    sequence = quicksortListByResID(chainBlock)
    blocksOfChains.append(sequence)

    listSortedByChainThenID = []   
    for block in blocksOfChains:
        listSortedByChainThenID.extend(block)
    
    return listSortedByChainThenID


def quicksortListByResID(itemList):

    if len(itemList) <= 1: return itemList
    
    pivot = itemList.pop()
    before = [item for item in itemList if item.resid <= pivot.resid]
    after = [item for item in itemList if item.resid > pivot.resid]
    
    
    return quicksortListByResID(before) + [pivot] + quicksortListByResID(after)    


def quicksortListByChain(itemList):
  
    if len(itemList) <= 1: return itemList
    
    pivot = itemList.pop()
    before = [item for item in itemList if item.chain <= pivot.chain]
    after = [item for item in itemList if item.chain > pivot.chain]
    
    
    return quicksortListByChain(before) + [pivot] + quicksortListByChain(after)

def getSequenceOfAlignment_web(site1, site2, outputPath, maxDist):

    pairsOfAtoms = matchAtoms(site1, site2, maxDist)
    reducedPairsOfAtoms = eliminateSideChains(pairsOfAtoms)

    atomsList = quicksortPairsByChainThenID(reducedPairsOfAtoms, 0)

    atomsListProtein = []
    atomsListDNAorRNA = []
    for atomPair in atomsList:
        if atomIsFromProtein(atomPair[0]):
            atomsListProtein.append(atomPair)
        elif atomIsFromDNAorRNA(atomPair[0]):
            atomsListDNAorRNA.append(atomPair)

    outFile = open(outputPath, 'w')

    if atomsListProtein:

        chainList1, idList1, aaList1, chainList2, idList2, aaList2, pipeList = getLinesOfSequenceAlignment_web(atomsListProtein)

        num = len(aaList1)

        outFile.write(('<tr><th colspan="%s"><b>Protein alignment</b></th></tr>\n') % (num+1))

        query_chain = '<tr><td><b>Chain:</b></td>'
        query_resid = '<tr><td><b>ResID:</b></td>'
        query_resname = '<tr><td><b>Query:</b></td>'

        pipe = '<tr><td> </td>'

        target_resname = '<td><b>Target:</b></td>'
        target_resid = '<td><b>ResID:</b></td>'
        target_chain = '<td><b>Chain:</b></td>'

        for i in xrange(num):

            query_chain += '<td>%s</td>' % (chainList1[i])
            query_resid += '<td>%s</td>' % (idList1[i])
            query_resname += '<td>%s</td>' % (aaList1[i])

            pipe += '<td>%s</td>' % (pipeList[i])

            target_resname += '<td>%s</td>' % (aaList2[i])
            target_resid += '<td>%s</td>' % (idList2[i])
            target_chain += '<td>%s</td>' % (chainList2[i])

        outFile.write(query_chain+'</tr>\n')
        outFile.write(query_resid+'</tr>\n')
        outFile.write(query_resname+'</tr>\n')

        outFile.write(pipe+'</tr>\n')

        outFile.write(target_resname+'</tr>\n')
        outFile.write(target_resid+'</tr>\n')
        outFile.write(target_chain+'</tr>\n')

    if atomsListDNAorRNA:

        chainList1, idList1, aaList1, chainList2, idList2, aaList2, pipeList = getLinesOfSequenceAlignment_web(atomsListDNAorRNA)

        num = len(aaList1)

        outFile.write(('<tr><th colspan="%s">DNA/RNA alignment</th></tr>\n') % (num+1))

        query_chain = '<tr><td><b>Chain:</b></td>'
        query_resid = '<tr><td><b>ResID:</b></td>'
        query_resname = '<tr><td><b>Query:</b></td>'

        pipe = '<tr><td> </td>'

        target_resname = '<td><b>Target:</b></td>'
        target_resid = '<td><b>ResID:</b></td>'
        target_chain = '<td><b>Chain:</b></td>'

        for i in xrange(len(aaList1)):

            query_chain += '<td>%s</td>' % (chainList1[i])
            query_resid += '<td>%s</td>' % (idList1[i])
            query_resname += '<td>%s</td>' % (aaList1[i])

            pipe += '<td>%s</td>' % (pipeList[i])

            target_resname += '<td>%s</td>' % (aaList2[i])
            target_resid += '<td>%s</td>' % (idList2[i])
            target_chain += '<td>%s</td>' % (chainList2[i])

        outFile.write(query_chain+'</tr>\n')
        outFile.write(query_resid+'</tr>\n')
        outFile.write(query_resname+'</tr>\n')

        outFile.write(pipe+'</tr>\n')

        outFile.write(target_resname+'</tr>\n')
        outFile.write(target_resid+'</tr>\n')
        outFile.write(target_chain+'</tr>\n')

    outFile.close()

def getLinesOfSequenceAlignment_web(atomsList):

    #
    atom1_1 = atomsList[0][0]
    chain1 = atom1_1.chain
    id1 = atom1_1.resid
    aa1 = getOneLetterResudueCode(atom1_1)
    #
    atom2_1 = atomsList[0][1]
    chain2 = atom2_1.chain
    id2 = atom2_1.resid
    aa2 = getOneLetterResudueCode(atom2_1)
    #

    #
    chainList1 = []
    chainList1.append('{0}'.format(chain1))
    idList1 = []
    idList1.append('{0}'.format(id1))
    aaList1 = []
    if atomIsFromLigand(atom1_1):
        aaList1.append('{0}{1}'.format(aa1, '*'))
    else:
        aaList1.append('{0}'.format(aa1))
    #
    pipeList = []
    #
    chainList2 = []
    chainList2.append('{0}'.format(chain2))
    idList2 = []
    idList2.append('{0}'.format(id2))
    aaList2 = []
    if atomIsFromLigand(atom2_1):
        aaList2.append('{0}{1}'.format(aa2, '*'))
    else:
        aaList2.append('{0}'.format(aa2))
    #

    if aaList1[0] == aaList2[0]:
        pipeList.append('|')
    else:
        pipeList.append(' ')

    k = 0
    for i in xrange(1, len(atomsList)):

        atom1_i  = atomsList[i][0]
        atom2_i  = atomsList[i][1]

        if (getOneLetterResudueCode(atom1_i) != None and getOneLetterResudueCode(atom2_i) != None):
            flag = False
            k=k+1
            # The first atom from a pair
            if atom1_i.chain != chain1:
                chainList1.append('|')
                idList1.append('|')
                aaList1.append('|')
                #
                pipeList.append(' ')
                #
                chainList2.append('|')
                idList2.append('|')
                aaList2.append('|')
                chainList1.append('{0}'.format(atom1_i.chain))
                chain1 = atom1_i.chain
                flag = True
                k=k+1
            else:
                chainList1.append(' ')

            if (atom1_i.resid != id1 + 1) and (atom1_i.resid != id1 - 1):
                if flag == False:
                    temp = chainList1[k]
                    chainList1[k] = ' '
                    chainList1.append(temp)
                    idList1.append('|')
                    aaList1.append('|')
                    #
                    pipeList.append(' ')
                    #
                    chainList2.append(' ')
                    idList2.append('|')
                    aaList2.append('|')
                    k=k+1
                    flag=True

                idList1[k-2] = '{0}'.format(atomsList[i-1][0].resid)
                idList1.append('{0}'.format(atom1_i.resid))
                id1 = atom1_i.resid
            else:
                idList1.append('{0}'.format(atom1_i.resid))
                id1 = atom1_i.resid

            if atomIsFromLigand(atom1_i):
                aaList1.append('{0}{1}'.format(getOneLetterResudueCode(atom1_i), '*'))
            else:
                aaList1.append('{0}'.format(getOneLetterResudueCode(atom1_i)))

            #The second atom from a pair
            if atom2_i.chain != chain2:
                #If chain is changed in second sequence, pipe is inserted in both sequenses and first sequence is shifted
                if flag == False:
                    temp = chainList1[k]
                    chainList1[k] = '|'
                    chainList1.append(temp)
                    temp = idList1[k]
                    idList1[k] = '|'
                    idList1.append(temp)
                    temp = aaList1[k]
                    aaList1[k] = '|'
                    aaList1.append(temp)
                    chainList2.append('|')
                    idList2.append('|')
                    aaList2.append('|')
                    pipeList.append(' ')
                    k=k+1
                    flag = True
                else:
                    chainList1[k-1] = '|'
                    chainList2[k-1] = '|'

                chainList2.append('{0}'.format(atom2_i.chain))
                chain2 = atom2_i.chain
            else:
                chainList2.append(' ')

            aaTemp = aaList1[k]
            if (atom2_i.resid != id2 + 1) and (atom2_i.resid != id2 - 1):
                if flag == False:
                    temp = chainList1[k]
                    chainList1[k] = '  '
                    chainList1.append(temp)
                    temp = idList1[k]
                    idList1[k] = '|'
                    idList1.append(temp)
                    temp = aaList1[k]
                    aaList1[k] = '|'
                    aaList1.append(temp)
                    temp = chainList2[k]
                    chainList2[k] = ' '
                    chainList2.append(temp)
                    idList2.append('|')
                    aaList2.append('|')
                    pipeList.append(' ')
                    k=k+1
                    flag = True

                idList2[k-2] = '{0}'.format(atomsList[i-1][1].resid)
                idList2.append('{0}'.format(atom2_i.resid))
                id2 = atom2_i.resid
            else:
                idList2.append('{0}'.format(atom2_i.resid))
                id2 = atom2_i.resid

            if atomIsFromLigand(atom2_i):
                aaList2.append('{0}{1}'.format(getOneLetterResudueCode(atom2_i), '*'))
            else:
                aaList2.append('{0}'.format(getOneLetterResudueCode(atom2_i)))

            if aaTemp == aaList2[-1]:
                pipeList.append('|')
            else:
                pipeList.append(' ')

    return chainList1, idList1, aaList1, chainList2, idList2, aaList2, pipeList

def get_atom_from_sequence(atom, sequence):
    
    atom_ = next((a for a in sequence if ((a != "-" )\
                                and (a.aa.strip().upper() == atom.aa.strip().upper())\
                                and (a.resid == atom.resid)\
                                and (a.chain.strip().upper() == atom.chain.strip().upper()))), None)  
    
    return atom_

def get_neighbor_atom(i, sequence_ref, sequence_old):
    
    neig_atom_old = None
    d = -1
    
    while neig_atom_old == None:
        d = d*(-1)
        q_ind = i-d
        if q_ind == len(sequence_ref):
            break
        if q_ind >= 0:
            neig_atom_link = sequence_ref[q_ind]
            neig_atom_old = get_atom_from_sequence(neig_atom_link, sequence_old)
        if d < 0:
            d -= 1
           
    return neig_atom_old, d

def match_signature_to_sequence(signature, sequence):
    
    matched_ligands = []
    mathed_positions = []
            
    ligands = re.split('X\(\d*\)', signature)
    indices = re.findall('X\(\d*\)', signature, flags=0)
    deltas = list(map(lambda x: int(x.strip('X(').strip(')')), indices))
    
    alignment = list(sequence)
            
    delta = 0
    i_stop = 0
    first_lig = ligands[0]
    for i in xrange(i_stop, len(alignment)):

        res = alignment[i]

        if res == first_lig:

            i_stop = i+1
            new_start = i+1
            matched_ligands.append(res)
            mathed_positions.append(i)
            
            for lg in xrange(1, len(ligands)):
                lig = ligands[lg]
                delta = deltas[lg-1]
                k=0
                for j in xrange(new_start, len(alignment)):
                    next_res = alignment[j]
                    if (next_res != '-'):
                        if (k == delta):
                            matched_ligands.append(alignment[j])
                            mathed_positions.append(j)
                            new_start = j+1
                            break
                        k += 1
            if matched_ligands == ligands:
                return matched_ligands, mathed_positions
            else:
                matched_ligands[:] = []
                mathed_positions[:] = []
                        
    return matched_ligands, mathed_positions

def fetch_nearby_residues(sequence, i_res, window):
    
    nearby_residues = [None for x in range(2*window)]
    
    n=1
    m=1
    while n <= window:
        
        ind_prv = i_res-m
        if ind_prv < 0:
            break
        res = sequence[ind_prv]
        if res != '-':
            ind = window - n
            nearby_residues[ind] = res
            n += 1
        m += 1

    n=1
    m=1
    while n <= window:
        
        ind_nxt = i_res+m
        if ind_nxt >= len(sequence):
            break
        res = sequence[ind_nxt]
        if res != '-':
            ind = window + n - 1
            nearby_residues[ind] = res
            n += 1
        m += 1
        
    return nearby_residues

def fetch_nearby_residues_old(sequence, i_res):
        
    prv_prv = ''
    prv = ''
    nxt = ''
    nxt_nxt = ''
    
    ind_prv_prv = ''
    ind_prv = ''
    ind_nxt = ''
    ind_nxt_nxt = ''

    if (i_res > 1):
        j=1
        flag_prv = True
        while sequence[i_res-j] == '-':
            j += 1
            if i_res-j == 0:
                flag_prv = False
                break
        if flag_prv:
            if i_res-j != 0:
                prv = sequence[i_res-j]
                ind_prv = i_res-j
                
                j_prv = j
              
                while sequence[i_res-j_prv-1] == '-':
                    j_prv += 1
                    if i_res-j_prv-1 == 0:
                        flag_prv = False
                        break
                if flag_prv:
                    prv_prv = sequence[i_res-j_prv-1]
                    ind_prv_prv = i_res-j_prv-1
    
    if (i_res < (len(sequence)-2)):
        j=1
        flag_nxt = True
        while sequence[i_res+j] == '-':
            j += 1
            if i_res+j == len(sequence):
                flag_nxt = False
                break
        if flag_nxt:
            nxt = sequence[i_res+j]
            ind_nxt = i_res+j
            
            j_nxt = j
            if i_res+j_nxt+1 != len(sequence):
                while sequence[i_res+j_nxt+1] == '-':
                    j_nxt += 1
                    if i_res+j_nxt+1 == len(sequence):
                        flag_nxt = False
                        break
                if flag_nxt:
                    nxt_nxt = sequence[i_res+j_nxt+1]
                    ind_nxt_nxt = i_res+j_nxt+1
    
    return [prv_prv, prv, nxt, nxt_nxt]

def get_ligand_proximity_pattern(ligand, ligand_proximity, window):
    
    lig_aa = getOneLetterResudueCode(ligand)
    ligand_proximity_aas = [getOneLetterResudueCode(x) if x !=None else x for x in ligand_proximity]
    
    ligand_proximity_aas.insert(window, lig_aa)
    pattern_residues = ['.' if x == None else x for x in ligand_proximity_aas]
    pattern = "{0}(-)*{1}(-)*{2}(-)*{3}(-)*{4}(-)*{5}(-)*{6}".format(*pattern_residues)
    
    return pattern

def get_ligands_mapping_to_sequence(ligands, ligands_proximity, sequence):
    
    import re
    
    ligands_map = []
    ligands_map_positions = []
    window = 3
    
    for l in xrange(len(ligands)):
            
        lig = ligands[l]
        lig_aa = getOneLetterResudueCode(lig)
        lig_proximity = ligands_proximity[l]
        
        if all(l is None for l in lig_proximity):
            continue
        
        pattern = get_ligand_proximity_pattern(lig, lig_proximity, window)
        
        r = re.search(pattern, sequence)
        if r:
            i=0
            m = r.span()[0]
            while i < window:
                res = sequence[m]
                if res == '-':
                    m += 1
                    continue
                i += 1
                m += 1
            
            ligands_map.append(sequence[m])
            ligands_map_positions.append(m)
    
    return ligands_map, ligands_map_positions