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

# basic PSSD output parser
# stores in Chain Residue format and also as elements

import os
import sys
import string
import traceback

from Log import LogFH
from Bio.PDB.PDBParser import PDBParser
from Bio.SCOP.Raf import to_one_letter_code
from GenericElement import GenericElement, GenericElementList

class Element:
    def __repr__(self):
        line = "< %c %c %4d %c %4d %c >" % (self.chain, self.type, self.start, self.starticode, self.end, self.endicode)
        return line

class PSSDResidue(dict):
    def __init__(self, logger, BioPDBRes, sstype):
        self.logger = logger
        self.BioPDBRes = BioPDBRes

        logger = self.logger
        res = self.BioPDBRes
        try:
            self.serial=None
            self.resid=res.id[1]
            self.icode=res.id[2]
            self.chain=res.get_parent().id
            self.aa=to_one_letter_code[res.resname]
            self.structure=sstype
            self.type=self.structure
            self.BP1=None
            self.BP2=None
            (x,y,z)=res['CA'].coord
            self.coord = res['CA'].coord
            self['CA'] = res['CA']
            if res.has_id('C'):
                self['C'] = res['C']
            if res.has_id('N'):
                self['N'] = res['N']
            if res.has_id('O'):
                self['O'] = res['O']
        except:
            [etype, val, tb] = sys.exc_info()
            errors = traceback.format_exception(etype, val, tb)
            logger.log(errors)

    def __repr__(self):
        line=''
        line+=self.chain+' '+str(self.resid)+self.icode+' '+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 pairBP(self):
        pass

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

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

class PSSD(dict):
    def __init__(self, logger=None):
        if logger:
            self.logger = logger
            
    def __sort(self, a, b):
        return cmp(self[self.chainID][a]['#'], self[self.chainID][b]['#'])

    def store(self, fname):
        if not os.path.isfile(fname): return 'Err'
        self.fname = os.path.basename(fname)
        fh=file(fname)
        lines=fh.readlines()
        fh.close()
        buf=[]
        for line in lines:
            if ( (line[:5] == 'HELIX') or
                 (line[:5] == 'SHEET') ):
                buf.append(line)
        sslines = buf
        self.oldelements = []
        for line in sslines:
            element = Element()
            if line[:5] == 'HELIX':
                element.chain = line[19]
                element.type = 'H'
                element.start = int(line[21:25])
                element.starticode = line[25]
                element.end = int(line[33:37])
                element.endicode = line[37]
            else:
                element.chain = line[21]
                element.type = 'E'
                element.start = int(line[22:26])
                element.starticode = line[26]
                element.end = int(line[33:37])
                element.endicode = line[37]
            self.oldelements.append(element)

        self.elements = PSSDElementList(self.logger)
        for oldelement in self.oldelements:
            element = PSSDElement(self.logger, oldelement.type)
            self.elements.append(element)

        parser = PDBParser(PERMISSIVE=1)
        try:
            structure = parser.get_structure('xxx', fname)
        except:
            self.logger.log('Bio.PDB could not parse: '+fname)
            return 'Err'
        self.structure = structure
        model = structure.get_list()[0]
        ctr = 0
        buf_elemlist = [False] * len(self.elements)
        for chain in model.get_list():
            self[chain.id] = {}
            structure = {'H':0, 'E':0}
            for res in chain.get_list():
                if not to_one_letter_code.has_key(res.resname):
                    continue
                ctr += 1
                aa = to_one_letter_code[res.resname]
                self[chain.id][(res.id[1], res.id[2])] = \
                                           {'#':ctr, 'AA':aa, 'SS':[]}
                residue = self[chain.id][(res.id[1], res.id[2])]
                for (i, elem) in enumerate(self.oldelements):
                    if ((chain.id == elem.chain) and \
                        (res.id[1] == elem.start ) and \
                        (res.id[2] == elem.starticode)):
                        structure[elem.type] += 1
                        buf_elemlist[i] = True
                        
                if structure['E']:
                    residue['SS'].append('E')
                if structure['H']:
                    residue['SS'].append('H')
                if not residue['SS']:
                    residue['SS'].append(' ')

                for (i, val) in enumerate(buf_elemlist):
                    if not val: continue
                    pssd_res = PSSDResidue(self.logger, \
                                           res, self.elements[i].type)
                    self.elements[i].append(pssd_res)

                for (i, elem) in enumerate(self.oldelements):
                    if ((chain.id == elem.chain) and \
                        (res.id[1] == elem.end ) and \
                        (res.id[2] == elem.endicode)):
                        structure[elem.type] -= 1
                        buf_elemlist[i] = False

    def show(self, outfh):
        self.chainID = self.keys()[0]
	residues = self[self.chainID].keys()
	residues.sort(self.__sort)
	for residue in residues:
	    print >> outfh, self[self.chainID][residue]['SS'][0]

if __name__ == '__main__':
    logger = LogFH('x')
    pssd = PSSD(logger)
    pssd.store(sys.argv[1])
    pssd.show(sys.stdout)
