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

from Bio.PDB.Polypeptide import CaPPBuilder
from Bio.SCOP.Raf import to_one_letter_code

import coordcalc
import Quad
import Ladder

class PSSDPeptides(list):
    def __init__(self, structure, parm, logger):
        self.parm = parm
        self.logger = logger
        #structure.build_CA_peptides(radius=5.0)
        #peptides = structure.get_CA_peptides()
  	ppb = CaPPBuilder(radius=5.0)
	peptides = ppb.build_peptides(structure)
        if not peptides:
            self.logger.log('critical', "No polypeptides found")
            return
        list.__init__(self, peptides)
        self.logger.log('info', "Polypeptides generated")
        self._find_residue_neighbours()
                
    def _find_residue_neighbours(self):
        '''
        mark residue neighbours (prev and next residues)
        mark residue serial number (within structure)
        res.prev[1] and res.next[1] are i-1 and i+1 ...
        res.prev[0] and res.next[0] are the same as res
        '''
        max_pos = 6
        max_pos += 1
        strpos = 0
        for peptide in self:
            length = len(peptide)
            for i in range(length):
                res = peptide[i]
                res.strpos = strpos
                strpos += 1
                res.prev = [None] * max_pos
                res.next = [None] * max_pos
                for j in range(max_pos):
                    if i >= j:
                        res.prev[j] = peptide[i-j]
                    if i < (length - j):
                        res.next[j] = peptide[i+j]
        self.logger.log('info', '_find_residue_neighbours() done')

    def _calc_d3(self):
        '''
        Loop over all residues in structure. Calculate d3
        '''
        for peptide in self:
            for res in peptide[:-3]:
                res.d3 = coordcalc.distance(res['CA'].coord, \
                                            res.next[3]['CA'].coord)
        self.logger.log('info', '_calc_d3() done')

    def _calc_dih(self):
        '''
        Loop over all residues in structure. Calculate dih
        '''
        for peptide in self:
            for res in peptide[:-3]:
                res.dih = coordcalc.torsion(res['CA'].coord, \
                                            res.next[1]['CA'].coord,
                                            res.next[2]['CA'].coord,
                                            res.next[3]['CA'].coord)
        self.logger.log('info', '_calc_dih() done')

    def _calc_i2ang(self):
        '''
        Loop over all residues in structure. Calculate i2angle
        '''
        for peptide in self:
            for res in peptide[2:-2]:
                v1 = res.prev[2]['CA'].coord - res['CA'].coord
                v2 = res['CA'].coord - res.next[2]['CA'].coord
                res.i2ang = coordcalc.v_angle(v1, v2, 180)
        self.logger.log('info', '_calc_i2ang() done')

    def _calc_i4avang(self):
        '''
        Loop over all residues in structure. Generate an average point
        with 4 consecutive residues. Find angle between residue before
        the 4 points, average point and residue after the 4 points.
        '''
        for peptide in self:
            for res in peptide[1:-4]:
                vCAcoord = (res['CA'].coord + \
                            res.next[1]['CA'].coord + \
                            res.next[2]['CA'].coord + \
                            res.next[3]['CA'].coord) / 4
                v1 = res.prev[1]['CA'].coord - vCAcoord
                v2 = vCAcoord - res.next[4]['CA'].coord
                res.i4avang = coordcalc.v_angle(v1, v2, 180)
        self.logger.log('info', '_calc_i4avang() done')

    def _calc_oh(self):
        '''
        Loop over all residues in structure. Calculate oh
        '''
        parm = self.parm
        for peptide in self:
            for res in peptide[:-3]:
                if ((res.d3 <= parm.hdist_cutoff) and \
                    (res.dih >= parm.hdih1_cutoff) and \
                    (res.dih <= parm.hdih2_cutoff )):
                    res.oh = 1
        self.logger.log('info', '_calc_oh() done')

    def _calc_os(self):
        '''
        Loop over all residues in structure. Calculate os
        '''
        parm = self.parm
        for peptide in self:
            for res in peptide[:-3]:
                if ((res.d3 > parm.sdist_strict_limit) and \
                    ((res.dih < parm.sdih1_cutoff) or \
                     (res.dih > parm.sdih2_cutoff))):
                    res.os = 1
        self.logger.log('info', '_calc_os() done')

    def _calc_ohh(self):
        '''
        Loop over all residues in structure. Calculate ohh
        '''
        parm = self.parm
        for peptide in self:
            for res in peptide[:-3]:
                if (res.d3 <= parm.hdist_strict_limit) and \
                    (res.dih >= (parm.hdih_optimal
                                     - (parm.hdih_sigma_range
                                        * parm.hdih_sigma))) and \
                    (res.dih <= (parm.hdih_optimal
                                     + (parm.hdih_sigma_range
                                        * parm.hdih_sigma))):
                    res.ohh = 1
        self.logger.log('info', '_calc_ohh() done')

    def _calc_h(self):
        '''
        Loop over all residues in structure. Calculate h
        '''
        for peptide in self:
            for res in peptide[:-4]:
                if res.oh and res.next[1].oh:
                    for j in range(5):
                        res.next[j].h = 1
        self.logger.log('info', '_calc_h() done')

    def _calc_s(self):
        '''
        Loop over all residues in structure. Calculate s
        '''
        for peptide in self:
            for res in peptide[:-3]:
                if res.os:
                    for j in range(4):
                        res.next[j].s = 1
        self.logger.log('info', '_calc_s() done')

    def _calc_hh(self):
        '''
        Loop over all residues in structure. Calculate hh
        '''
        for peptide in self:
            for res in peptide[:-5]:
                if res.ohh and \
                       res.next[1].ohh and \
                       res.next[2].ohh:
                    for j in range(6):
                        res.next[j].hh = 1
        self.logger.log('info', '_calc_hh() done')

    def _insert_bulge(self):
        for peptide in self:
            for res in peptide[1:-1]:
                if res.prev[1].E and \
                       res.next[1].E and \
                       not res.E:
                    if res.next[1] in res.prev[1].pairs['all']: continue
                    prevpairs = res.prev[1].pairs['all']
                    nextpairs = res.next[1].pairs['all']
                    linked = 0
                    for i in prevpairs:
                        for j in nextpairs:
                            if ((j==i.prev[1]) or \
                                (j==i.next[1]) or \
                                (j==i.prev[2] and i.prev[1].pairs['all']) or \
                                (j==i.next[2] and i.next[1].pairs['all'])):
                                res.E = 1
                                res.blg = [None]
                                m='PSSDPeptides._insert_bulge() '
                                m+=res.get_parent().id+str(res.id)
                                self.logger.log('trace',m)
                                linked = 1
                            if linked: break
                        if linked: break
        self.logger.log('info', '_insert_bulge() done')

    def _explain(self):
        lines = []
        lines.append('Format: Column: 1\tchain\n')
        lines.append('Format: Column: 2-5\tresid\n')
        lines.append('Format: Column: 6\ticode\n')
        lines.append('Format: Column: 8\taltloc\n')
        lines.append('Format: Column: 10-12\tresname\n')
        lines.append('Format: Column: 14\tresname (one letter code)\n')
        lines.append('Format: Column: 34\thelix tag\n')
        lines.append('Format: Column: 45\tstrand tag\n')        
        lines.append('Format: Column: 56\tdssp\n')
        lines.append('Format: Column: 58\tstride\n')
        lines.append('Format: Column: 60\tsstruc\n')
        lines.append('Format: Column: 62\tpross\n')
        lines.append('Format: Column: 64\tmolauto\n')        
        lines.append('Format: Column: 66-73\td3\n')
        lines.append('Format: Column: 75-82\ti,i1,i2,i3 torsion\n')
        lines.append('Format: Column: 84-91\ti-2,i,i2 angle\n')
        lines.append('Format: Column: 93-100\ti-1,av(i,i1,i2,i3),i4 angle\n')
        lines.append('Format: Column: 102-112\tstrand pairs\n')
        return lines

    ## Public Methods

    def defout(self):
        ATOM_FORMAT_STRING="%s%5i %-4s%c%3s %c%4i%c   %8.3f%8.3f%8.3f%6.2f%6.2f      %4s%2s%2s\n"

	def _get_atom_line(atom, hetfield, segid, atom_number, resname, \
                           resseq, icode, chain_id, element="  ", charge="  "):
            if hetfield!=" ":
                record_type="HETATM"
            else:
                record_type="ATOM  "
            name=atom.get_fullname()
            altloc=atom.get_altloc()
            x, y, z=atom.get_coord()
            bfactor=atom.get_bfactor()
            occupancy=atom.get_occupancy()
            args=(record_type, atom_number, name, altloc, resname, chain_id, \
                  resseq, icode, x, y, z, occupancy, bfactor, segid, \
                  element, charge)
            return ATOM_FORMAT_STRING % args

        lines = []
        atom_number = 0
        for peptide in self:
            chain_id = peptide[0].get_parent().get_id()
            for res in peptide:
                atom = res['CA']
                hetfield, resseq, icode = res.get_id()
                segid = res.get_segid()
                atom_number += 1
                resname = res.get_resname()
                line = _get_atom_line(atom, hetfield, segid, atom_number, \
                                      resname, resseq, icode, chain_id)
                line = line[:77]+line[13]+line[78:]
                if res.H:
                    line = line[:67]+'H'+line[68:]
                if res.E:
                    line = line[:69]+'E'+line[70:]
                lines += [line]
        return lines

    def log(self):
        lines = self.defout()
        self.logger.log('peptides', lines)
        lines = []
        lines += self._explain()
        for peptide in self:
            chain = peptide[0].get_parent()
            for residue in peptide:
                if to_one_letter_code.has_key(residue.resname):
                    rescode = to_one_letter_code[residue.resname]
                else: rescode = ' '
                line = '%c%4d%c %c %s %c ' % \
                       (chain.get_id(), residue.id[1],
                        residue.id[2], residue['CA'].altloc,
                        residue.resname, rescode)
                taglist = ['oh','os','ohh','s','h','hh']
                taglist += ['H','q1','q2','q3','E','blg','hpin']
                for attr in taglist:
                    if hasattr(residue, attr) and \
                           getattr(residue, attr):
                        line += ' '+attr
                    else:
                        space = ' '
                        line += space*(len(attr)+1)
                taglist = ['compare_dssp','compare_stride','compare_sstruc']
                taglist += ['compare_pross', 'compare_molauto']
                for attr in taglist:
                    if hasattr(residue, attr):
                        val = getattr(residue, attr)
                        if val:
                            line += ' '+val
                        else:
                            line += '  '
                taglist = ['d3','dih','i2ang','i4avang']
                for attr in taglist:
                    if hasattr(residue, attr) and \
                           getattr(residue, attr):
                        line += ' %8.3f' % getattr(residue, attr)
                    else:
                        line += '         '
                if hasattr(residue, 'pairs'):
                    for p in residue.pairs['all']:
                        line += ' '+p.get_parent().id+'%4d'%p.id[1]
                line += '\n'
                lines += [line]
        self.logger.log('peptides', lines)

    def initialise_tags(self, taglist, ValFunc=None):
        if not ValFunc:
            def ValFunc():
                return None
        for peptide in self:
            for residue in peptide:
                for tag in taglist:
                    setattr(residue, tag, ValFunc())

    def generate_values(self):
        taglist = ['d3','dih','i2ang','i4avang']
        self.initialise_tags(taglist)
        self._calc_d3()
        self._calc_dih()
        self._calc_i2ang()
        self._calc_i4avang()

    def interpret_values(self):
        taglist = ['oh','os','ohh','h','s','hh']
        self.initialise_tags(taglist)
        self._calc_oh()
        self._calc_os()
        self._calc_ohh()
        self._calc_h()
        self._calc_s()
        self._calc_hh()

    def run_primary_ruleset(self):
        self.generate_values()
        self.interpret_values()
        self.logger.log('info', "_run_primary_ruleset() done")

    def build_quadruplets(self):
        ladders = Ladder.PSSDLadders(self.parm, self.logger)
        def add_quad(resi, resj, orientation):
            if resi.hh and resi.next[1].hh: return
            if resj.hh and resj.next[1].hh: return
            dist = [None] * 2
            if orientation == 'parallel':
                dist[0] = coordcalc.dist_sqr(resi['CA'].coord,
                                             resj['CA'].coord)
                dist[1] = coordcalc.dist_sqr(resi.next[1]['CA'].coord,
                                             resj.next[1]['CA'].coord)
            elif orientation == 'antiparallel':
                dist[0] = coordcalc.dist_sqr(resi['CA'].coord,
                                             resj.next[1]['CA'].coord)
                dist[1] = coordcalc.dist_sqr(resi.next[1]['CA'].coord,
                                             resj['CA'].coord)
            for d in dist:
                if d > self.parm.max_sheet_pair_cutoff_sqr: return
            quad = Quad.Quad(self.parm, self.logger, level=1)
            quad.build(resi, resj, orientation)
            if quad.valid['q1']:
                quad.log()
                ladders.append(quad)
        for ppi in range(len(self)):
            pepi = self[ppi]
            pepi_len = len(pepi)
            for pepj in self[ppi:]:
                if pepi == pepj:
                    for i in range(1,pepi_len-8):
                        for j in range(i+6,pepi_len-2):
                            add_quad(pepi[i], pepj[j], 'parallel')
                    for i in range(1,pepi_len-6):
                        for j in range(i+4,pepi_len-2):
                            add_quad(pepi[i], pepj[j], 'antiparallel')
                else:
                    pepj_len = len(pepj)
                    for i in range(1,pepi_len-2):
                        for j in range(1,pepj_len-2):
                            add_quad(pepi[i], pepj[j], 'parallel')
                            add_quad(pepi[i], pepj[j], 'antiparallel')
        ladders.sort(Quad.QuadCmp)
        self.logger.log('info', '_quad_gen() done '+str(len(ladders))+' quads')
        return ladders
   
    def and_tag(self, source, target):
        for peptide in self:
            for residue in peptide:
                if getattr(residue, source):
                    setattr(residue, target, 1) 

    def mark_strands(self):
        self.initialise_tags('E')
        self.and_tag('q1', 'E')
        self.and_tag('q2', 'E')
        self.and_tag('q3', 'E')
        self._insert_bulge()
        self.logger.log('info', 'mark_strands() done')

    def delinvalidrulemarks(self, dict):
        for peptide in self:
            for res in peptide[1:-1]:
                for (type, list) in dict.items():
                    for tag in list:
                        if not (hasattr(res, tag) and getattr(res, tag)):
                            continue
                        if hasattr(res.prev[1], type) and \
                           getattr(res.prev[1], type) and \
                           hasattr(res, type) and \
                           getattr(res, type) and \
                           hasattr(res.next[1], type) and \
                           getattr(res.next[1], type):
                            continue
                        delattr(res, tag)

