# $CVSHeader: pal8e/src/DSSPElements.py,v 1.6 2005/02/15 01:10:25 indraneel Rel $

# considers key (list)value pairing for element type and residue type
# does _not_ consider left handed Helices
# reads dssp file and converts to elements (lists of residues)

import sys
import os
import traceback
from string import split
from Log import LogFH
from types import *
from Numeric import array, Float0
from GenericElement import GenericElement, GenericElementList

class DSSPType:
    def __init__(self, logger, type_definition):
        self.type_definition = type_definition

    def getType(self, DSSPStructure):
        if not self.type_definition: return DSSPStructure
        if self.type_definition.has_key(DSSPStructure):
            return self.type_definition[DSSPStructure]
        return ' '

class DSSPResidue:
    def __init__(self, logger, line, mol, dsspdict):
        self.logger = logger
        self.line = line
        self.mol = mol
        self.DSSPDict = dsspdict

    def __repr__(self):
        line=''
        line+=self.chain+' '+str(self.resid)+' '+self.structure
        for attr in ['pair1', 'pair2']:
            if hasattr(self, attr):
                line+=' '+ '%4d' % (getattr(self,attr)).resid
            else: line+='     '
        line+=' '+self.type
        line+=' '+str(self.coord)
        return line

    def mkResidue(self, typedef):
        logger = self.logger
        try:
            self.serial=int(split(self.line[1:5])[0])
            self.resid=int(split(self.line[6:10])[0])
            self.icode=self.line[10]
            self.chain=self.line[11]
            self.aa=self.line[13]
            self.structure=self.line[16]
            self.handedness=self.line[22]
            if (self.structure == 'H') and (self.handedness == '-'):
                self.structure = ' '
            self.type=typedef.getType(self.structure)
            self.BP1=int(split(self.line[25:29])[0])
            self.BP2=int(split(self.line[29:33])[0])
            self.phi=float(self.line[103:109])
            self.psi=float(self.line[109:115])
            x=float(self.line[115:122])
            y=float(self.line[122:129])
            z=float(self.line[129:136])
            self.coord = array((x, y, z), Float0)
            if not self.mol.has_key(self.chain):
                self.mol[self.chain]={}
            self.mol[self.chain][self.resid]=self
            self.DSSPDict[self.serial] = self
        except:
            [etype, val, tb] = sys.exc_info()
            errors = traceback.format_exception(etype, val, tb)
            logger.log([self.line])
            logger.log(errors)

    def pairBP(self):
        DSSPDict = self.DSSPDict
        self.pairs = []
        if self.BP1:
            self.bp1 = DSSPDict[self.BP1].resid
            self.pair1 = DSSPDict[self.BP1]
            self.pairs.append(self.pair1)
        else:
            self.bp1 = 0
        if self.BP2:
            self.bp2 = DSSPDict[self.BP2].resid
            self.pair2 = DSSPDict[self.BP2]
            self.pairs.append(self.pair2)
        else:
            self.bp2 = 0

class DSSPElement(GenericElement):
    def __init__(self, logger, type, seq=None):
        GenericElement.__init__(self, logger, type, seq)

class DSSPElementList(GenericElementList):
    def __init__(self, logger, seq=None):
        GenericElementList.__init__(self, logger, seq)

class DSSPmkElements:
    def __init__(self, logger, lines, dsspdict):
        self.logger = logger
        self.lines = lines
        self.DSSPDict = dsspdict

    def getElements(self, logger, typedef):
        elements = DSSPElementList(logger)
        mol = {}
        oldresidue = None
        newresidue = None
        oldelement = None
        newelement = None
        for line in self.lines:
            if line[13] == '!':
                newresidue = None
            else:
                newresidue = DSSPResidue(logger, line, mol, self.DSSPDict)
                newresidue.mkResidue(typedef)
            if (oldresidue and (
                (line[13]=='!') or
                (oldresidue.type != newresidue.type) or
                (oldresidue.chain != newresidue.chain)
                )):
                if newelement:
                    oldelement = newelement
                    elements.append(oldelement)
                    newelement = None
            if newresidue:
                if not newelement:
                    newelement = DSSPElement(logger, newresidue.type)
                newelement.append(newresidue)
                oldresidue = newresidue
                newresidue = None
        if oldresidue:
            oldelement = newelement
            elements.append(oldelement)
        elements.pairBP()
        self.elements = elements
        return self.elements

class DSSPExtract:
    def __init__(self, logger, lines):
        self.logger = logger
        self.lines = lines

    def getLines(self):
        lines = self.lines
        for (i,line) in zip(range(len(lines)+1),lines):
            if line[2] == '#':
                break
        lines = lines[i+1:]
        return lines

class DSSPElements:
    def __init__(self, logger, DSSPfh_or_fname):
        self.DSSPDict = {}
        self.logger = logger
        if type(DSSPfh_or_fname) == StringType:
            self.fh = file(DSSPfh_or_fname)
            self.lines = self.fh.readlines()
            self.fh.close()
        else:
            self.fh = DSSPfh_or_fname
            self.lines = self.fh.readlines()

    def setType(self, type_definition):
        self.type_definition = DSSPType(self.logger, type_definition)

    def getElements(self):
        logger = self.logger
        lines = self.lines
        extractor = DSSPExtract(logger, lines)
        lines = extractor.getLines()
        element_maker = DSSPmkElements(logger, lines, self.DSSPDict)
        typedef = self.type_definition
        elements = element_maker.getElements(logger, typedef)
        return elements

if __name__ == '__main__':
    logger = LogFH('DSSPElements-testlog')
    dsspdir = '../../dat/nrchn2787.dssp/'
    for fName in os.listdir(dsspdir):
        fname = dsspdir + fName
        dssp_elements = DSSPElements(logger, fname)
        typedef={'H':'H', 'G':'H', 'I':'I', 'E':'E'}
        dssp_elements.setType(typedef)
        elements = dssp_elements.getElements()
        #for elem in elements:
        #    if elem.type == 'I':
        #        print fname
        #        print elem
