# $CVSHeader: pal8e/src/Ladder.py,v 1.8 2005/02/15 01:10:26 indraneel Rel $

import Quad
from CoordCalc import dist_sqr
from CoordCalc import v_angle

class PSSDLadder(list):
    def __init__(self, parm, logger):
        self.parm = parm
        self.logger = logger

    def __repr__(self):
        line = ''
        for q in self:
            line += str(q)+'\n'
        return line

    def _findedge(self, side, quad):
        edge = {}
        if side == 'b':
            edge['side'] = 'b'
            edge[1] = quad.b1
            edge[2] = quad.b2
            edge['drift1'] = 'prev'
            edge['revdrift1'] = 'next'            
            if quad.orient == 'parallel':
                edge['drift2'] = 'prev'
                edge['revdrift2'] = 'next'                
            else:
                edge['drift2'] = 'next'
                edge['revdrift2'] = 'prev'                
            edge['quad'] = quad
        elif side == 'c':
            edge['side'] = 'c'
            edge[1] = quad.c1
            edge[2] = quad.c2
            edge['drift1'] = 'next'
            edge['revdrift1'] = 'prev'            
            if quad.orient == 'parallel':
                edge['drift2'] = 'next'
                edge['revdrift2'] = 'prev'                
            else:
                edge['drift2'] = 'prev'
                edge['revdrift2'] = 'next'                
            edge['quad'] = quad
        return edge

    def _chkangle(self, res1, res2, res3):
        angle = v_angle((res1['CA'].coord - res2['CA'].coord), \
                        (res3['CA'].coord - res2['CA'].coord))
        if angle < self.parm.sheet_plane_angle:
            #print res1, res2, res3
            return None
        return 'allow'

    def _climb(self, side, quad, qtype, disallowed_reslist):
        '''
        propagate the edge if possible. The edge residues are already paired.
        '''
        max_cutoff_dist = self.parm.max_sheet_pair_cutoff
        min_cutoff_dist = self.parm.min_sheet_pair_cutoff
        max_cutoff_dist_sqr = max_cutoff_dist * max_cutoff_dist
        min_cutoff_dist_sqr = min_cutoff_dist * min_cutoff_dist

        newquads = []
        while 1:
            edge = self._findedge(side, quad)
            drift1 = edge['drift1']
            drift2 = edge['drift2']
            revdrift1 = edge['revdrift1']
            revdrift2 = edge['revdrift2']
            edge1 = getattr(edge[1], drift1)[1]
            edge2 = getattr(edge[2], drift2)[1]
            
            if not (edge1 and edge2): break
            if len(edge1.pairs['all']) > 1: break
            if len(edge2.pairs['all']) > 1: break
            if edge2 in edge1.pairs['all']: break
            if edge1.hh or edge2.hh: break
            if not (edge1.s and edge2.s): break
            if edge1 in disallowed_reslist: break
            if edge2 in disallowed_reslist: break
            if edge1 in [edge2, edge2.next[1], edge2.next[2]]: break
            dist_square = dist_sqr(edge1['CA'].coord, edge2['CA'].coord)
            if dist_square > max_cutoff_dist_sqr: break
            if dist_square < min_cutoff_dist_sqr: break

            brk = 0
            for (res, revdrift) in [(edge1, revdrift1), (edge2, revdrift2)]:
                revres = getattr(res, revdrift)[3]
                if not (revres.q1 or revres.q2): continue
                cutoff = self.parm.sdist_strict_limit
                if (revdrift == 'next') and (res.d3 < cutoff):
                    brk = 1
                elif (revdrift == 'prev') and \
                         (revres.d3 < self.parm.sdist_strict_limit):
                    brk = 1
                if brk: break
            if brk: break
            
            if edge1.pairs['all']:
                allow = self._chkangle(edge1.pairs['all'][0], edge1, edge2)
                if not allow: break
            if edge2.pairs['all']:
                allow = self._chkangle(edge2.pairs['all'][0], edge2, edge1)
                if not allow: break
            if edge1.hpin:
                allow = self._chkangle(edge1.hpin, edge1, edge2)
                if not allow: break
            if edge2.hpin:
                allow = self._chkangle(edge2.hpin, edge2, edge1)
                if not allow: break

            a1 = getattr(edge1, revdrift1)[2]
            a2 = getattr(edge2, revdrift2)[2]
            b1 = getattr(edge1, revdrift1)[1]
            b2 = getattr(edge2, revdrift2)[1]
            c1 = edge1
            c2 = edge2
            d1 = getattr(edge1, drift1)[1]
            d2 = getattr(edge2, drift2)[1]
            if edge1 == edge2.prev[1]:
                d1 = None
                d2 = None
            prevquad = quad
            quad = Quad.Quad(self.parm, self.logger)
            if drift1 == drift2:
                orient = 'parallel'
            else: orient = 'antiparallel'
            quad.generate(a1, b1, c1, d1, a2, b2, c2, d2, orient)
            if not quad.valid[qtype]: break
            if len(self) == 1:
                for res in edge['quad'].res:
                    if len(res.pairs['all']) < 2: continue
                    for qtype in ['q2', 'q3']:
                        if not res.pairs[qtype]: continue
                        for pres in res.pairs[qtype]:
                            if pres in edge['quad'].pair[res].prev[:2] or \
                                   pres in edge['quad'].pair[res].next[:2]:
                                pres = res.pairs[qtype][0]
                                pres.pairs[qtype].remove(res)
                                pres.pairs['all'].remove(res)
                                pres.blg.append(res)
                                res.pairs[qtype].remove(pres)
                                res.pairs['all'].remove(pres)
                                break
                edge['quad'].pairup('q1')
            if edge1 == edge2.prev[1]:
		edge1.hpin = edge2
		edge2.hpin = edge1
            quad.pairup(qtype)
            prevquad.link[side] = quad
            if side == 'b':
                quad.link['c'] = prevquad
                self.insert(0, quad)
            elif side == 'c':
                quad.link['b'] = prevquad
                self.append(quad)
            newquads.append(quad)
            quad.log()
            if edge1 in edge2.prev[1:3]: break
            if qtype == 'q3' and len(self) > 1: break
        return newquads

    ## Public Methods

    def extend(self, qtype, disallowed_reslist):
        newquads = []
        for (side, quad) in [('b', self[0]), ('c', self[-1])]:
            if qtype == 'q3' and len(self) > 1: break
            newquads += self._climb(side, quad, qtype, disallowed_reslist)
        return newquads

## PSSDLadders is interpreted from disjoint randomly linked lists
class PSSDLadders(list):
    def __init__(self, parm, logger):
        self.parm = parm
        self.logger = logger

    def _link_quads(self):
        '''
        Link quadruplets
        (Simply link quadruplet to another quadruplet to generate sheets.
        Go down the list, selecting quadruplets one at a time.)
        For every quadruplet, check if it can fit a quad in the list
        above it. Do not check for a fit if the quad above has both edges
        linked. Do not link anti/parallel quads.
        Every time there is a link discard quads (including references)
        in the list below which cannot fit.
        Bulges are marked on residues which come from a worse quadruplet.
        '''
        for quad in self:
            for _quad in self:
                if (quad.link['b'] and quad.link['c']) or \
                       (quad == _quad):
                    ## delete alternatives
                    for q in self:
                        if q.score <= quad.score:
                            continue
                        if quad.is_similar_to(q):
                            self.remove(q)
                    break
                ## check if free edges fit
                _quad.checkandfit(quad)
        self.logger.log('info', '_link_quads() done')

    def _remove_double_bulge_quads(self):
        delbuf = []
        for quad in self:
            ctr = 0
            for res in quad.res:
                if quad.pair[res] in res.blg:
                    ctr += 1
            if ctr > 1:
                delbuf.append(quad)
        for quad in delbuf:
            quad.isolate()
            self.remove(quad)
        self.logger.log('info', '_remove_double_bulge_quads() done')

    def _quad_discard(self, pairpool, failctr):
        for r, i in failctr.items():
            if i == 2:
                for quad in pairpool[r]:
                    quad.delete = 1

    def _find_ladder_length(self, quad):
        length = 1
        for s in ['b', 'c']:
            q = quad
            while q.link[s]:
                q = q.link[s]
                if hasattr(q, 'delete'): break
                length += 1
        return length

    def _quad_choose(self, pairpool, failctr):
        allpassed = 0
        if 1 not in failctr.values():
            if len(pairpool) == 2:
                return
            if len(pairpool) == 3:
                allpassed = 1
                for r in pairpool.keys():
                    failctr[r] = 1
        lenpool = {}
        for r, i in failctr.items():
            if i != 1: continue
            length = self._find_ladder_length(pairpool[r][0])
            if lenpool.has_key(length):
                prevscores = []
                for q in lenpool[length]:
                    prevscores.append(q.score)
                prevscores.sort()
                newscores = []
                for q in pairpool[r]:
                    newscores.append(q.score)
                newscores.sort()
                if prevscores[0] < newscores[0]:
                    for q in pairpool[r]:
                        q.delete = 1
                    continue
                else:
                    for q in lenpool[length]:
                        q.delete = 1
            lenpool[length] = pairpool[r]
        klist = lenpool.keys()
        klist.sort()
        delist = klist[:-1]
        if allpassed: delist = klist[:-2]
        for k in delist:
            for quad in lenpool[k]:
                quad.delete = 1

    def _quadpaircheck(self, res, pairpool, itr):
        '''
        side has the common residues in the 3 quads making up the quadpool
        '''
        if len(pairpool) == 2:
            positions = [(0,1)]
        else:
            positions = [(0,1), (1,2), (2,0)]
        pairres = pairpool.keys()
        failctr = {}
        for r in pairres:
            failctr[r] = 0
        for i,j in positions:
            paira = pairres[i]
            pairb = pairres[j]
            angle = v_angle((paira['CA'].coord - res['CA'].coord), \
                            (pairb['CA'].coord - res['CA'].coord))
            if angle < self.parm.sheet_plane_angle:
                failctr[paira] += 1
                failctr[pairb] += 1
        if itr == 'discard':
            self._quad_discard(pairpool, failctr)
        elif itr == 'choose':
            self._quad_choose(pairpool, failctr)

    def _quadfilter(self, itr):
        '''
        Do not allow quads that do not pass pair angle cutoff
        with any other quad.
        '''
        for quad in self:
            if hasattr(quad, 'delete'): continue
            for res in quad.res:
                pairpool = {}
                for quad in self:
                    if hasattr(quad, 'delete'): continue
                    if res in quad.res:
                        pair = quad.pair[res]
                        if res in pair.blg or \
                           pair in res.blg:
                            continue
                        if not pairpool.has_key(pair):
                            pairpool[pair] = [quad]
                        else:
                            pairpool[pair].append(quad)
                num = len(pairpool)
                if num < 2: continue
                if num > 3:
                    line = '_quadfilter() '+itr+' : '
                    line += '>3 pairs for '+res.get_parent().id
                    line += str(res.id[1])+res.id[2]
                    self.logger.log('critical', line)
                if (itr == 'discard') and (num != 3): continue
                self._quadpaircheck(res, pairpool, itr)
        delbuf = []
        for quad in self:
            if hasattr(quad, 'delete'):
                delbuf.append(quad)
                continue
        for quad in delbuf:
            quad.isolate()
            self.remove(quad)
        self.logger.log('info', '_quadfilter('+itr+') done')

    def _pairup_quads(self, single=0):
        for quad in self:
            if not single and not (quad.link['b'] or quad.link['c']):
                continue
            quad.pairup('q1')
        if not single:
            self.logger.log('info', 'Residues paired')

    def _remove_lone_quads(self):
        delbuf = []
        for quad in self:
            if quad.link['b'] or quad.link['c']: continue
            if quad.b1.q1 and quad.c1.q1:
                if quad.orient == 'parallel':
                    prev = quad.b2.prev[1]
                    next = quad.c2.next[1]
                else:
                    prev = quad.b2.next[1]
                    next = quad.c2.prev[1]
                if not (prev.H and next.H):
                    continue
            if quad.b2.q1 and quad.c2.q1 and not \
               (quad.b1.prev[1].H and quad.c1.next[1].H):
                continue
            delbuf.append(quad)
        for quad in delbuf:
            quad.isolate()
            self.remove(quad)
        self.logger.log('info', 'Lone quadruplets removed')

    def _getladderlist(self):
        for quad in self:
            quad.tag = None
        ladderlist = []
        for quad in self:
            if quad.link['b']: continue
            if quad.tag: continue
            ladder = PSSDLadder(self.parm, self.logger)
            q = quad
            while q:
                ladder.append(q)
                q.tag = 1
                q = q.link['c']
            ladderlist.append(ladder)
        return ladderlist

    ## Public Methods

    def initiate(self, peptides):
        peptides.initialise_tags(["blg"], list)
        self._link_quads()
        self._remove_double_bulge_quads()
        self._quadfilter('discard')
        self._quadfilter('choose')
        def valfunc():
            return {'all':[], 'q1':[], 'q2':[], 'q3':[]}
        peptides.initialise_tags(['pairs'], valfunc)
        peptides.initialise_tags(['q1'])
        self._pairup_quads()
        self._remove_lone_quads()
        self.logger.log('info', 'Ladders initiated')

    def extend(self, peptides, definition):
        disallowed_reslist = []
        for elem in definition:
            if elem.type == 'H':
                disallowed_reslist += elem[2:-2]
        peptides.initialise_tags(['q2', 'q3', 'hpin'])
        ladderlist = self._getladderlist()
        def _cmp(a, b):
            return cmp(len(a), len(b))
        ladderlist.sort(_cmp)
        ladderlist.reverse()
        for qtype in ['q2', 'q3']:
            for ladder in ladderlist:
                self += ladder.extend(qtype, disallowed_reslist)
        self.logger.log('info', "Ladders extended")
        
