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

from Element import PSSDElement, PSSDElements
import helixbrk
from StatMaker import HelixbrkStat, StrandbrkStat, ElemDistOvlapStat

class PSSDefinition(PSSDElements):
    def __init__(self, parm, logger):
        self.parm = parm
        self.logger = logger
        PSSDElements.__init__(self, logger)

    def _grow_helices(self, type='H'):
        for elem in self:
            if elem.type != type: continue
            lastres = elem.end()
            elem += lastres.next[1:4]

    def _prune_helix(self, elem):
        '''
        No residue of helix can have 2 pairs
        Max overlap with strand at edge is 2 residue
        5 Residue helix cannot have overlap on both sides with strand
        '''
        if len(elem) < 5: return
        # check from middle of helix
        size = len(elem)
        if float(size/2) == float(size)/2:
            midi = size/2
        else: midi = size/2 + 1
        delbuf = []
        for i in range(midi-1, -1, -1):
            res = elem[i]
            if len(res.pairs['all']) > 1:
                m = res.get_parent().id
                m += str(res.id)
                self.logger.log('trace', m, 'helix pruned >2 pairs')
                delbuf += elem[:i+1]
                break
            if res.pairs['all'] and res.next[1].pairs['all']:
                delbuf += elem[:i]
                break
        for i in range(midi+1, len(elem)):
            res = elem[i]
            if len(res.pairs['all']) > 1:
                m = res.get_parent().id
                m += str(res.id)
                self.logger.log('trace', m, 'helix pruned >2 pairs')
                delbuf += elem[i:]
                break
            if res.pairs['all'] and res.prev[1].pairs['all']:
                delbuf += elem[i+1:]
                break
        for res in delbuf:
            elem.remove(res)
        if len(elem) < 5: return
        if (len(elem) == 5):
            if elem.start().E and \
               elem.end().E:
                x = elem.pop(0)
                x = elem.pop()

    def _helix_has_ohh(self, elem):
        for res in elem[:-3]:
            if res.ohh: return 1
        return None
        
    def _verify_helices(self, use_strands=0):
        '''
        Remove <5 residue helix
        Remove helix with no 'ohh' residue
        '''
        delbuf = []
        for i in range(len(self)):
            elem = self[i]
            if elem.type != 'H': continue
            if use_strands:
                self._prune_helix(elem)
            if len(elem) < 5:
                delbuf.append(i)
                continue
            if not self._helix_has_ohh(elem):
                delbuf.append(i)
                continue
        for i in range(len(delbuf)):
            j = delbuf[i] - i
            self.pop(j)
        for elem in self:
            if elem.type == 'H':
                elem.tag_residues('H')
        self.logger.log('info', '_verify_helices() done')

    def _addbroken(self, buf, elements, type):
        '''
        adds list of broken elements (list of lists of residues)
        in proper format
        '''
        for seq in elements:
            element = PSSDElement(type, self.logger, seq)
            buf.append(element)

    def _helixsplit(self, helices, brkres):
        pieces = []
        overlaps = None
        for helix in helices:
            buf = []
            for res in helix:
                buf.append(res)
                if res == brkres:
                    buf += res.next[1:3]
                    overlaps = res.next[1:3]
                    pieces.append(buf)
                    buf = []
            if buf: pieces.append(buf)
        for helix in pieces:
            if len(helix) < 5: break
        else:
            if overlaps:
                for res in overlaps:
                    res.Hr1o = 1
            return pieces
        return helices

    def _helixsplitbyaxis(self, helix):
        '''
        input a helix (list of residues)
        output broken helices (list of lists of residues)
        this uses rotational fit using RMSD along helix axis
        '''
        data = []
        if len(helix) >= 8:
            for resi in helix:
                for j in resi['CA'].coord:
                    data.append(j)
            buf = helixbrk.splithelix(data);
        else: buf = None
        if (buf == None) or (len(buf) == 1):
            return [helix]
        brks = {}
        for i in buf[1:]:
            lbrks = {}
            for j in range(i['start']-5,i['start']+5):
                if not helix[j].ohh:
                    lbrks[helix[j].d3] = helix[j]
            for (d3, res) in lbrks.items():
                if res not in helix[2:-4]:
                    del lbrks[d3]
            if not lbrks: continue
            ky = lbrks.keys()
            ky.sort()
            brks[ky[-1]] = lbrks[ky[-1]]
        ky = brks.keys()
        ky.sort()
        ky.reverse()
        helices = [helix]
        for i in ky:
            helices = self._helixsplit(helices, brks[i])
        return helices

    def _splithelices(self):
        '''
        Break already defined helices
        '''
        #self.statmaker = HelixbrkStat(self.logger) ## stat
        newhelices = []
        for elem in self:
            if elem.type != 'H': continue
            #if (len(elem) >= 8):
                #self.statmaker.stathelix_recursive(helix)
                #self.statmaker.stathelix_incremental(helix)
            helices = self._helixsplitbyaxis(elem)
            #if len(helices) > 1:
                #print elem
                #for h in helices:
                    #print h
                #print
            self._addbroken(newhelices, helices, 'H')
        #self.statmaker.dumphelixstat_recursive()  ## stat
        #self.statmaker.dumphelixstat_incremental() ## stat
        delbuf = []
        for i in range(len(self)):
            if self[i].type == 'H':
                delbuf.append(i)
        for i in range(len(delbuf)):
            j = delbuf[i] - i
            self.pop(j)
        for elem in newhelices:
            self.append(elem)
        for elem in self:
            if elem.type == 'H':
                elem.tag_residues('H')
                #print elem
                #print elem.getEnds()
                #print elem.getRMSD()
                #rmsd = elem.getRMSD()
                #if (rmsd > 2.5) and (len(elem) > 12):
                    #print rmsd, elem

    ## find if list1 and list2 contain nearly consecutive residues
    def _has_consecutive(self, list1, list2, gap=2):
        for i in list1:
            for j in list2:
                if (i in j.prev[1:gap+1]) or \
                       (i in j.next[1:gap+1]):
                    return 'has_consecutive'
        return None

    ## find if res in strand1 or strand2
    def _hairpin_res_places(self, reslist):
        places = [None] * len(reslist)
        for (i,res) in enumerate(reslist):
            if ((i != 0) and (places[i-1] == 2)):
                places[i] = 2
                continue
            if not res.pairs['all']:
                places[i] = 0
                continue
            if self._has_consecutive(res.pairs['all'], \
                                     res.prev[1].pairs['all'], 1):
                places[i] = 1
                continue
            elif self._has_consecutive(res.pairs['all'], \
                                       res.next[1].pairs['all'], 1):
                places[i] = 2
        for (i,res) in enumerate(reslist):
            if places[i]: continue
            if ((i != 0) and (places[i-1] == 2)):
                places[i] = 2
                continue
            if not res.pairs['all']: continue
            if self._has_consecutive(res.pairs['all'], \
                                     res.prev[1].pairs['all'], 2):
                places[i] = 1
                continue
            elif self._has_consecutive(res.pairs['all'], \
                                       res.next[1].pairs['all'], 2):
                places[i] = 2
        return places

    def _strandbreak_hairpin(self, strand):
        strands = []
        buf = []
        skiplist = []
        for res in strand:
            if res in skiplist: continue
            buf.append(res)
            if res in strand[-3:]: continue
            if res.next[1] == res.hpin:
                strands.append(buf)
                buf = []
                continue
            if res == strand[-4]: continue
            if res.next[2] in res.pairs['all']:
                places = self._hairpin_res_places(res.next[1:2])
                if places[0] == 0:
                    skiplist.append(res.next[1])
                    res.next[1].blg=[]
                elif places[0] == 1:
                    buf.append(res.next[1])
                    skiplist.append(res.next[1])
                strands.append(buf)
                buf = []
                continue
            if res.next[3] in res.pairs['all']:
                res.next[1].hpin = res.next[2]
                res.next[2].hpin = res.next[1]
                buf.append(res.next[1])
                strands.append(buf)
                buf = []
                skiplist.append(res.next[1])
                continue
            if res == strand[-5]: continue
            if res.next[4] in res.pairs['all']:
                places = self._hairpin_res_places(res.next[1:4])
                for (i,place) in zip([1,2,3], places):
                    if place == 0:
                        skiplist.append(res.next[i])
                        res.next[i].blg = []
                    elif place == 1:
                        buf.append(res.next[i])
                        skiplist.append(res.next[i])
                strands.append(buf)
                buf = []
                continue
        if buf:
            strands.append(buf)
        return strands
        
    def _strandbreakd3(self, strand):
        '''
        i and i+3 cannot be in same strand if d3 < 8
	and no bulges exist on the strand. previous strand ends at i+2
        '''
        strands = []
        buf = []
        do_not_brk_lst = []
        for res in strand:
            buf.append(res)
            if res in do_not_brk_lst: continue
            if res in strand[-3:]: continue
            if res.d3 > self.parm.sdist_strict_limit: continue
            for r in res.next[0:4]:
	        if hasattr(r, 'blg') and getattr(r, 'blg'): break
	    else:
                buf += res.next[1:3]
                strands.append(buf)
                buf = []
                do_not_brk_lst = res.next[1:3]
        if buf:
            strands.append(buf)
        num = len(strands)
        if num == 1: return strands
        for i in range(1, num):
            if strands[i][0] in strands[i-1][-2:]:
                strands[i][0].Er1o = 1
                if strands[i][1] == strands[i-1][-1]:
                    strands[i][1].Er1o = 1
            else:
                strands[i][0].Er1s = 1
        return strands

    def _strandbreak45(self, strand):
        '''
        Break a strand when it bends (i2ang > 45) and
        is not a bulge (i4avang > 70)
        '''
        strands = []
        buf = []
        do_not_brk_lst = []
        for res in strand:
            buf.append(res)
            if res in do_not_brk_lst: continue
            if res in strand[:2]: continue
            if res in strand[-4:]: continue
            if (res.i4avang > self.parm.sheet_bulge_angle) and \
                   (res.i2ang > self.parm.sheet_bend_angle):
                #print res
                buf.append(res.next[1])
                strands.append(buf)
                buf = []
                do_not_brk_lst = [res.next[1]]
        if buf:
            strands.append(buf)
        num = len(strands)
        if num == 1: return strands
        for i in range(1, num):
            if strands[i][0] in strands[i-1][-2:]:
                strands[i][0].Er2o = 1
                if strands[i][1] == strands[i-1][-1]:
                    strands[i][1].Er2o = 1
            else:
                strands[i][0].Er2s = 1
        return strands

    def _strandbreak_pairing(self, strand):
        '''
        Break if consecutive residues do not link to consecutive residues.
        Check for bulges.
        case 1:          case 2:           
          o = = o          o = = o         
          o = = o          o = = o         
          o = = o          o = = o          <== i0
                o = = o    o = = o = = o    <== i1
                o = = o          o = = o    <== i2
                o = = o          o = = o
        exception (bulge):
          o = = o = = o    o = = o = = o    o = = o = = o
          o = = o = = o    o = = o = = o    o = = o = = o
          o = = o                o                o
          o = = o = = o    o = = o = = o    o = = o
          o = = o = = o    o = = o = = o    o = = o
        '''
        strands = []
        buf = []
        for res in strand:
            buf.append(res)
            if res == strand[-1]: continue
            pairs = [[], [], []]
            for i in range(3):
                if res.next[i]:
                    pairs[i] += res.next[i].pairs['all']
                    if res.next[i].hpin:
                        pairs[i].append(res.next[i].hpin)
            if (not res.blg) and (not res.next[1].blg) and \
                   not self._has_consecutive(pairs[0], pairs[1]):
                strands.append(buf)
                buf = []
                continue
            if res == strand[-2]: continue
            if self._has_consecutive(pairs[0], pairs[1]) and \
                   self._has_consecutive(pairs[1], pairs[2]) and \
                   (not res.blg) and (not res.next[2].blg) and \
                   not self._has_consecutive(pairs[0], pairs[2], gap=3):
                buf.append(res.next[1])
                strands.append(buf)
                buf = []
        if buf: strands.append(buf)
        return strands

    def _splitstrands(self):
        #self.statmaker = StrandbrkStat(self.logger) ## stat
        fnlist = []
        fnlist += [self._strandbreak_hairpin]
        fnlist += [self._strandbreakd3]
        fnlist += [self._strandbreak45]
        fnlist += [self._strandbreak_pairing]
        origstrands = []
        for elem in self:
            if elem.type != 'E': continue
            origstrands.append(elem)
            #self.statmaker.statstrand_i2_i4(elem)
        #self.statmaker.dumpstrandstat_i2_i4()  ## stat
        for fn in fnlist:
            newstrands = []
            for elem in origstrands:
                if len(elem) > 3:
                    strands = fn(elem)
                else:
                    strands = [elem]
                self._addbroken(newstrands, strands, 'E')
            origstrands = newstrands
        delbuf = []
        for i in range(len(self)):
            if self[i].type == 'E':
                delbuf.append(i)
        for i in range(len(delbuf)):
            j = delbuf[i] - i
            self.pop(j)
        for elem in newstrands:
            self.append(elem)
        self.logger.log('info', '_split_strands() done')

    def _get_complete_overlaps(self):
        self.sort_by_position()
        strands = []
        for i in range(len(self)):
            elem = self[i]
            if elem.type != 'E': continue
            elem.tag = i
            strands.append(elem)
        overlaps = []
        for (i, strand) in enumerate(strands):
            if len(strand) > 4: continue
            respool = []
            if i != 0:
                respool += strands[i-1]
            if i != (len(strands)-1):
                respool += strands[i+1]
            ctr = 0
            for res in strand:
                if res in respool:
                    ctr += 1
                    continue
                if not res.pairs['all']:
                    ctr += 1
            if ctr == len(strand):
                overlaps.append(strand.tag)
        def _cmp(a, b):
            return cmp(len(self[a]), len(self[b]))
        overlaps.sort(_cmp)
        return overlaps

    def _remove_complete_overlaps(self):
        while 1:
            overlaps = self._get_complete_overlaps()
            if not overlaps: break
            self.pop(overlaps[0])

    def _drop_strand(self, strand):
        if not strand in self: return
	for res in strand:
	    for pair in res.pairs['all']:
	        if res in pair.pairs['all']:
		    pair.pairs['all'].remove(res)
	    res.blg = []
	    for q1pair in res.pairs['q1']:
	        if res in q1pair.pairs['q1']:
		    q1pair.pairs['q1'].remove(res)
	    for q2pair in res.pairs['q2']:
	        if res in q2pair.pairs['q2']:
		    q2pair.pairs['q2'].remove(res)
	self.remove(strand)

    def _verify_strands(self):
        '''
        remove strands without at least 3 non helix-overlaping residues.
	remove strands with <2 consecutive residues from good quads.
        '''
        overlaps = []
        singles = []
	xtnd_only = []
	nopairs = []
        for elem in self:
            if elem.type != 'E': continue
            ctr = 0
            for res in elem:
                if res.H: ctr += 1
            if len(elem) - ctr < 3:
                if ctr == 0:
                    pass
                #singles.append(i)
                elif elem[0].H and elem[-1].H:
                    overlaps.append(elem)
	    for res in elem:
	        if res == elem[-1]: continue
		if res.q1 and res.next[1].q1: break
	    else:
	        xtnd_only.append(elem)
	    for res in elem:
	        if not (res.pairs['all'] or res.blg or res.hpin):
		    nopairs.append(elem)
	for i in [overlaps, singles, xtnd_only, nopairs]:
	    for elem in i:
	        self._drop_strand(elem)
        for elem in self:
            if elem.type != 'E': continue
            elem.tag_residues('E')
        self.logger.log('info', '_verify_strands() done')
    #

    ## Public Methods

    def initiate_helix(self, peptides):
        self.buildnew(['oh'], peptides)
        self.change_type('oh', 'H')
        self._grow_helices()
        peptides.initialise_tags(['H'])
        self._verify_helices()
        self.logger.log('info', '_initiate_helix() done')
        
    def remove_hairpin_helix(self, peptides):
        delbuf = []
        for i in range(len(self)):
            elem = self[i]
            if elem.type != 'H': continue
            if len(elem) > 7: continue
            hpin = 0
            for resi in elem:
                for resj in elem:
                    if resi in resj.pairs['all']:
                        hpin = 1
                        break
                if hpin: break
            if hpin:
                ctr = 0
                for res in elem:
                    if res.H and not res.q1: ctr += 1
                if ctr < 5:
                    delbuf.append(i)
        for i in range(len(delbuf)):
            j = delbuf[i] - i
            self.pop(j)
        peptides.initialise_tags(['H'])
        for elem in self:
            if elem.type == 'H':
                elem.tag_residues('H')
        self.logger.log('info', 'Removed hairpin helices')

    def get_E_pairs(self):
        def _cmp(resa, resb):
            return cmp(resa.strpos, resb.strpos)
        pairs = {}
        for elem in self:
            if elem.type != 'E': continue
            for res in elem:
                if res.pairs['all']:
                    if not pairs.has_key(res):
                        pairs[res] = []
                    for pair in res.pairs['all']:
                        if _cmp(res, pair) < 0:
                            if not pair in pairs[res]:
                                pairs[res].append(pair)
        for key in pairs.keys():
            pairs[key].sort(_cmp)
        return pairs

    def break_helices(self, peptides):
        self._verify_helices(use_strands=1)
        peptides.initialise_tags(['H'])
        self._splithelices()
        self.logger.log('info', 'break_helices() done')

    def break_strands(self, peptides):
        self._splitstrands()
        peptides.initialise_tags(['E'])
        self._remove_complete_overlaps()
        # temporary
        #self._verify_strands()
        for elem in self:
            if elem.type != 'E': continue
            elem.tag_residues('E')
        # temporary
        self.logger.log('info', 'break_strands() done')
            
    def buildoverlap(self, type1, type2, tag):
        for elemi in self:
            if elemi.type != type1: continue
            for elemj in self:
                if elemi == elemj: continue
                if elemj.type != type2: continue
                for res in elemi:
                    if res in elemj:
                        setattr(res, tag, 1)

    def buildstart(self, type, tag, count):
        for elem in self:
            if elem.type != type: continue
            for res in elem[:count]:
                setattr(res, tag, 1)

    def rollback_non_alpha_helices(self, peptides):
        self.buildnew(['oh'], peptides)
        self._grow_helices('oh')
        delbuf = []
        for i in range(len(self)):
            elem = self[i]
            if elem.type != 'oh': continue
            if len(elem) < 6:
                delbuf.append(i)
                continue
            brk = 0
            if self._helix_has_ohh(elem):
                delbuf.append(i)
                continue
            ctr = 0
            for res in elem:
                if res.H or res.E:
                    ctr = 0
                else:
                    ctr += 1
                    if ctr == 5 : break
            else:
                delbuf.append(i)
                continue
            # 5 consecutive (non H) (non E) located
            self._prune_helix(elem)
            if len(elem) < 5:
                delbuf.append(i)
        for i in range(len(delbuf)):
            j = delbuf[i] - i
            self.pop(j)
        self.change_type('oh', 'H')        
        for elem in self:
            if elem.type == 'H':
                elem.tag_residues('H')

    def filter_length(self, typelist, length):
        delbuf = []
        for i in range(len(self)):
            elem = self[i]
            if elem.type not in typelist: continue
            if len(elem) >= length:
                continue
            delbuf.append(i)
        for i in range(len(delbuf)):
            j = delbuf[i] - i
            self.pop(j)

    def report_badstrands(self):
        for (i,elem) in enumerate(self):
            if elem.type != 'E': continue
            if len(elem) >5: continue
            if len(elem.pairs)==0:
                self.logger.log('trace', str(elem),
                                'report_badstrands pairs=0')
                continue
            if len(elem.pairs)==1: continue
            if len(elem.pairs)>2:
                self.logger.log('trace', str(elem),
                                'report_badstrands pairs>2')
                for pair in elem.pairs:
                    self.logger.log('trace', str(pair),
                                    'report_badstrands pairs>2 (pair)')
                continue
            if elem.hasConsecutivePairs():
                self.logger.log('trace', str(elem),
                                'report_badstrands pairs=2')
                for pair in elem.pairs:
                    self.logger.log('trace', str(pair),
                                    'report_badstrands pairs=2 (pair)')

    def stat(self):
        #self.statmaker = ElemDistOvlapStat(self.logger) ## stat
        self.sort_by_position()
        for (i, elem1) in enumerate(self):
            if ((elem1.type != 'H') and (elem1.type != 'E')):
                continue
            for (j, elem2) in enumerate(self):
                if (j<=i): continue
                if ((elem2.type != 'H') and (elem2.type != 'E')):
                    continue
                if ((elem1.type == 'E') and (len(elem1) < 3)): continue
                if ((elem2.type == 'E') and (len(elem2) < 3)): continue
                if ((elem1.type == 'H') and (len(elem1) < 5)): continue
                if ((elem2.type == 'H') and (len(elem2) < 5)): continue
                #self.statmaker.statElem(elem1, elem2)
        #self.statmaker.dumpelemstat()
