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

from Scientific.Geometry import VectorModule
from CoordCalc import v_angle, dist_sqr
import Numeric
import math
import sys
import nick

class PSSDElement(list):
    def __init__(self, type, logger, seq=None):
        self.type = type
        self.logger = logger
        if seq:
            list.__init__(self, seq)

    def __repr__(self):
        line = 'ElementType: '+ self.type+'    '
        line += 'Chain: ' + self[0].get_parent().id + '    '
        line += '%4d' % (self[0].id[1])
        line += self[0].id[2] + ' <====> '
        line += '%4d' % (self[-1].id[1])
        line += self[-1].id[2]
        return line

    def _getEnds(self, start, end):
        v = VectorModule.Vector(end - start)
        u = v/v.length()
        # get projection of end points on unit vector
        pstart = self[0]['CA'].coord
        vstart = VectorModule.Vector(pstart-start)
        projStart = vstart * u
        pend = self[-1]['CA'].coord
        vend = VectorModule.Vector(pend - end)
        projEnd = vend * u
        # get start of element on vector
        estart = start + (projStart * u).array
        eend = end + (projEnd * u).array
        return (estart, eend)
    
    # returns endpoints calculated by rotFit or any other
    def _getEndsH(self, fn=None):
        if fn == None:
            fn = nick.rotFitPy
        data = []
        for res in self:
            for coord in res['CA'].coord:
                data.append(coord)
        #(start, end) = nick.rotFitPy(data, 3)
        (start, end) = fn(data, 3)
        start = Numeric.array(start)
        end = Numeric.array(end)
        return self._getEnds(start, end)

    def _getEndsE(self):
        leng = len(self)
        start = self[int(leng*0.3)]['CA'].coord
        end = self[int(leng*0.6)]['CA'].coord
        return self._getEnds(start, end)

    def _minmaxrmsd(self, lengths, overlaps):
        helices = []
        for (i,lg) in enumerate(lengths):
            # create an element (set the length)
            start = sum(lengths[:i]) - sum(overlaps[:i+1])
            seq = self[start:start+lg]
            element = PSSDElement('H', self.logger, seq)
            helices.append(element)
        rmsds = [0] * len(helices)
        for (i,h) in enumerate(helices):
            rmsds[i] = h.getRMSD()
        rmsds.sort()
        if (rmsds[-1] < self.maxrmsd):
            self.maxrmsd = rmsds[-1]
            self.helices = helices

    def _recursebrks(self, lengths, overlaps, brks):
        lenself = len(self)
        if (brks == -1):
            allocated = sum(lengths) - sum(overlaps)
            if (allocated != lenself): return self.helices
            self._minmaxrmsd(lengths, overlaps)
            return self.helices
            
        for ovlap in range(3):
            overlaps[brks] = ovlap
            if ((brks == 0) and (ovlap != 0)): continue
            #overlaps[0] = 0
            for hlen in range(5, lenself+1):
                lengths[brks] = hlen
                self._recursebrks(lengths, overlaps, brks - 1)
                
    def start(self):
        return self[0]

    def end(self):
        return self[-1]

    def log(self, handle=None):
        line = self.__repr__()
        self.logger.log('element', line, handle)

    def get_H_line(self, helix_ctr):
        line = ''
        line += "HELIX "+" "+ "%3d" % (helix_ctr) +" "+"   "+" "
        line += self[0].get_resname() + " "
        line += self[0].get_parent().get_id() + " "
        line += "%4d" % (self[0].get_id()[1])
        line += self[0].get_id()[2] + " "
        line += self[-1].get_resname() + " "
        line += self[-1].get_parent().get_id() + " "
        line += "%4d" % (self[-1].get_id()[1])
        line += self[-1].get_id()[2]
        line += "%2d" % (1)
        line += "                              " + " "
        line += "%5d" % (len(self))
        line += '\n'
        return line

    def get_E_line(self):
        line = ''
        line += "SHEET "+" "
        line += "   "+" "
        if hasattr(self, 'sheetID'):
            sid = '%3s' % self.sheetID
            line += sid
        else: line += '   '
        line += "  "+" "
        line += self[0].get_resname() + " "
        line += self[0].get_parent().get_id()
        line += "%4d" % (self[0].get_id()[1])
        line += self[0].get_id()[2] + " "
        line += self[-1].get_resname() + " "
        line += self[-1].get_parent().get_id()
        line += "%4d" % (self[-1].get_id()[1])
        line += self[-1].get_id()[2]
        line += "                                " + " "
        line += "%5d" % (len(self))
        line += '\n'
        return line

    def get_S_line(self):
        line = ''
        line += "EXTND " + " "
        line += '\n'

    def tag_residues(self, tag, value=1):
        # tagholders must exist with initialise_tags()
        for res in self:
            setattr(res, tag, value)

    def pairs_with(self, other):
        links = 0
        for res in self:
            pairs = res.pairs['all'] + [res.hpin]
            for pair in pairs:
                if pair == None: continue
                if pair in other:
                    links += 1
        return links

    def link_back(self):
        for res in self:
            res.element = self

    def link_back_pairs(self):
        for res in self:
            res.elempairs = []
            for pair in res.pairs['all']:
                if hasattr(pair, 'element') and \
                       (pair.element in self.pairs):
                    res.elempairs.append(pair.element)

    def hasConsecutivePairs(self):
        pair0 = self.pairs[0]
        pair1 = self.pairs[1]
        if (pair0[0].get_parent() != pair1[0].get_parent()): return False
        for res0 in pair0:
            for res1 in pair1:
                if (abs(res0.id[1] - res1.id[1]) < 3): return True
        return False

    def getEnds(self, fn=None):
        if ((len(self) >=5 ) and (self.type == 'H')):
            return self._getEndsH(fn)
        if ((len(self) >=3 ) and (self.type == 'E')):
            return self._getEndsE()

    def getRMSD(self):
        (pointP, pointQ) = self.getEnds()
        unitVector = VectorModule.Vector(pointQ-pointP)
        unitVector = unitVector / unitVector.length()
        SD = 0
        for R in self:
            vecPR = R['CA'].coord - pointP
            vecPR = VectorModule.Vector(vecPR)
            vecD = vecPR.cross(unitVector)
            vecSD = vecD * vecD
            SD += vecSD
        MSD = SD / len(self)
        RMSD = math.sqrt ( MSD )
        return RMSD

    def getAngle(self, elem):
        (a, b) = self.getEnds()
        (c, d) = elem.getEnds()
        v1 = b-a
        v2 = d-c
        angle = v_angle(v1, v2)
        return angle

    # will recurse to find best break points using RMSD only
    # does not use zee or angle, max brkpoints need to be specified
    # will not iterate over breakpoints unless specified
    def getBrokenHelices(self, brks, iterate=False):
        self.maxrmsd = 'NotSet'
        self.helices = []
        for b in range(brks+1):
            if ((not iterate) and (b != brks)): continue
            lengths = [0] * (b+1)
            overlaps = [0] * (b+1)
            self._recursebrks(lengths, overlaps, b)
        return self.helices

    def getDist(self, elem):
        minimum = sys.maxint
        for res in self:
            for r in elem:
                dsqr = dist_sqr(res['CA'].coord, r['CA'].coord)
                if (dsqr < minimum): minimum = dsqr
        minimum = math.sqrt(minimum)
        return minimum

    # proj[0] and proj[1] from self
    # proj[2] and proj[3] from elem (2 cases with reversed element)
    def getOverlap(self, elem):
        (sst, sen) = self.getEnds()
        (es, ee) = elem.getEnds()
        overlaps = []
        for (est, een) in [[es, ee], [ee, es]]:
            midPst = (sst + est)/2.0
            midPen = (sen + een)/2.0
            v = VectorModule.Vector(midPen - midPst)
            u = v/v.length()
            proj = []
            for p in [sst, sen, est, een]:
                v = VectorModule.Vector(p - midPst)
                proj.append(v * u)
            if ((proj[0] == proj[1]) or
                (proj[2] == proj[3])):
                overlaps.append(0.0)
                continue
            if (proj[0] <= proj[2] <= proj[1]):
                if (proj[3] <= proj[0]):
                    overlap = proj[2] - proj[0]
                elif (proj[3] >= proj[1]):
                    overlap = proj[1] - proj[2]
                else:
                    overlap = abs(proj[3] - proj[2])
            elif (proj[2] < proj[0]):
                if (proj[3] <= proj[0]):
                    overlap = 0
                elif (proj[3] >= proj[1]):
                    overlap = proj[1] - proj[0]
                else:
                    overlap = proj[3] - proj[0]
            else: # proj[2] > proj[1]
                if (proj[3] <= proj[0]):
                    overlap = proj[1] - proj[0]
                elif (proj[3] - proj[1]):
                    overlap = 0
                else:
                    overlap = proj[1] - proj[3]
            overlaps.append(overlap)
        overlap = max(overlaps)
        return overlap
            
class PSSDElements(list):
    def __init__(self, logger):
        self.logger = logger

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

    def log(self, handle=None):
        typelist = []
        for element in self:
            if element.type not in typelist:
                typelist.append(element.type)
        for type in typelist:
            for element in self:
                if element.type == type:
                    element.log(handle)

    def buildnew(self, types, peptides):
        for type in types:
            # clear existing elements
            ## do not change (list.remove(list) does not work)
            delbuf = []
            for i in range(len(self)):
                if self[i].type == type:
                    delbuf.append(i)
            for i in range(len(delbuf)):
                j = delbuf[i] - i
                self.pop(j)
            # build new elements
            elem = PSSDElement(type, self.logger)
            for peptide in peptides:
                for residue in peptide:
                    if hasattr(residue, type) and getattr(residue, type):
                        elem.append(residue)
                    elif elem:
                        self.append(elem)
                        elem = PSSDElement(type, self.logger)
                if elem:
                    self.append(elem)
                    elem = PSSDElement(type, self.logger)

    def change_type(self, oldtype, newtype):
        for element in self:
            if element.type == oldtype:
                element.type = newtype

    def sort_by_length(self):
        def _cmp(a, b):
            return cmp(len(a), len(b))
        self.sort(_cmp)

    def sort_by_position(self):
        def _cmp(a, b):
            return cmp(a.start().strpos, b.start().strpos)
        self.sort(_cmp)

    def pairup_strands(self):
        self.sort_by_position()
        for (i, elem) in zip(range(len(self)), self):
            elem.serial = i
            if elem.type == 'E':
                elem.pairs = []
        for strandi in self:
            if strandi.type != 'E': continue
            for strandj in self:
                if strandj.type != 'E': continue
                if strandi.serial >= strandj.serial: continue
                links = strandi.pairs_with(strandj)
                if links > 2:
                    if strandj not in strandi.pairs:
                        strandi.pairs.append(strandj)
                    if strandi not in strandj.pairs:
                        strandj.pairs.append(strandi)
        def _cmp(a,b):
            return cmp(len(a), len(b))
        for strand in self:
            del strand.serial
            if strand.type != 'E': continue
            if strand.pairs:
                strand.pairs.sort(_cmp)
                strand.pairs.reverse()

    def link_back(self):
        for element in self:
            element.link_back()

    def link_back_pairs(self):
        for element in self:
            element.link_back_pairs()
