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

# parses P-Sea output files and stores for comparisons
# stores in dict format

import os, sys
import string
from Log import LogFH

from Bio.PDB.PDBParser import PDBParser
from Bio.SCOP.Raf import to_one_letter_code

class PSEA(dict):
    def __init__(self, logger):
        self.logger = logger

    def __sort(self, a, b):
        return cmp(self[self.chainID][a]['#'], self[self.chainID][b]['#'])

    def setPDBDir(self, pdbDir, pdbExt='.pdb'):
        self.pdbDir = pdbDir
        self.pdbExt = pdbExt

    def setChain(self, chainID):
        self.chainID = chainID
            
    def store(self, fname):
        if not os.path.isfile(fname): return 'Err'
        self.fname = os.path.basename(fname)
        fileID = string.join(self.fname.split('.')[:-1], '')
        pdbFname = self.pdbDir+fileID+self.pdbExt
        if not os.path.isfile(pdbFname):
            return 'Err'
        fh = file(fname)
        psealines = fh.readlines()
        fh.close()
        for (i, line) in enumerate(psealines):
            if line[:6] == '>p-sea': break
        psealines = psealines[i+1:]
        psealines = map(lambda x: x[:-1], psealines)
        pseaSS = string.join(psealines, '')
        parser = PDBParser(PERMISSIVE=1)
        try:
            structure = parser.get_structure('xxx', pdbFname)
        except:
            self.logger.log('Bio.PDB could not parse: '+pdbFname)
            return 'Err'
        self.structure = structure
        model = structure.get_list()[0]
        ctr = 0
        #print pseaSS, len(pseaSS)
        try:
            for chain in model.get_list():
                if self.chainID != chain.id: continue
                self[chain.id] = {}
                #print len(chain.get_list())
                for res in chain.get_list():
                    if not to_one_letter_code.has_key(res.resname):
                        continue
                    if not res.has_id('CA'): continue
                    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])]
                    while pseaSS[ctr] == '-': ctr += 1
                    residue['SS'] = [pseaSS[ctr]]
                    ctr += 1
        except:
            self.logger.log('PSEA store() failed on '+self.fname)
            raise
            return 'Err'

    def getSortedRes(self):
        residues = self[self.chainID].keys()
        residues.sort(self.__sort)
        return residues

if __name__ == '__main__':
    logger = LogFH('x')
    psea = PSEA(logger)
    psea.setPDBDir('../../dat/cull100/', '.pdb')
    fname = sys.argv[1]
    chainID = fname[-5]
    psea.setChain(chainID)
    psea.store(fname)
    res = psea.getSortedRes()
    #for r in res:
    #    print psea[chainID][r]['SS'][0]
    logger.close()
