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

# gets d3 and dih from dssp for residues in pssd elements

import os, sys
import string
from Bio.PDB.PDBParser import PDBParser

from Log import LogFH
from Peptides import PSSDPeptides
from Compare_DSSP import DSSP_Reader
from Element import PSSDElements

logfile = '../../tmp/statlog'
statfile = '../../pssd-stat/dssp_d3_dih'
pdbdir = '../../dat/nrchn2787/'
dsspdir = '../../dat/nrchn2787.dsspout/'
dssp_extension = '.dssp'
taglist = ['H', 'G', 'I', 'E', 'B']
vallist = ['d3', 'dih']

class StatCollector:
    def __init__(self, taglist, vallist):
        self.taglist = taglist
        self.reslines = []
        self.counter = {}
        for tag in taglist:
            self.counter[tag] = {}
            for val in vallist:
                self.counter[tag][val] = {}

    def _get_resline(self, tag, res, chainid, structureid):
        line = ''
        line += tag + '\t'
        line += ('%.2f' % res.d3) + '\t'
        line += ('%.2f' % res.dih) + '\t'
        line += structureid + '\t'
        line += chainid + '\t'
        line += str(res.id[1]) + '\t'
        line += res.id[2]
        return line

    def _get_statlines(self):
        lines = []
        for tag in self.taglist:
            for val in vallist:
                keys = self.counter[tag][val].keys()
                keys.sort()
                for key in keys:
                    line = ''
                    line += tag+'\t'
                    line += val+'\t'
                    line += key+'\t'
                    line += str(self.counter[tag][val][key])
                    lines.append(line+'\n')
        return lines

    def collect(self, tag, res, chainid, structureid):
        resline = self._get_resline(tag, res, chainid, structureid)
        self.reslines.append(resline+'\n')
        ctr = self.counter[tag]
        d3 = '%.1f' % res.d3
        if ctr['d3'].has_key(d3):
            ctr['d3'][d3] += 1
        else:
            ctr['d3'][d3] = 1
        dih = str(int(round(res.dih)))
        if ctr['dih'].has_key(dih):
            ctr['dih'][dih] += 1
        else:
            ctr['dih'][dih] = 1

    def show(self):
        if os.path.isfile(statfile):
            os.remove(statfile)
        statlines = self._get_statlines()
        fh = file(statfile, 'a')
        fh.writelines(self.reslines)
        fh.write('---- total ----\n')
        fh.writelines(statlines)
        fh.close()

class SSMap:
    def __init__(self, structure, secstrdef, logger):
        self.structure = structure
        self.secstrdef = secstrdef
        self.logger = logger
        peptides = PSSDPeptides(structure, None, logger)
        self.peptides = peptides

    def _map(self):
        self.peptides.initialise_tags(taglist)
        model = self.structure[0]
        secstrdef = self.secstrdef
        chains = secstrdef.keys()
        for chain in chains:
            reslist = secstrdef[chain].keys()
            for resid in reslist:
                icodes = secstrdef[chain][resid].keys()
                for icode in icodes:
                    secstr = secstrdef[chain][resid][icode]
                    if secstr not in taglist: continue
                    res = (' ', resid, icode)
                    setattr(model[chain][res], secstr, 1)

    def run(self, stats):
        self._map()
        self.peptides.generate_values()
        elements = PSSDElements(self.logger)
        elements.buildnew(taglist, self.peptides)
        for element in elements:
            if (element.type == 'H') and (len(element) < 6): continue
            #if (element.type == 'E') and (len(element) < 4): continue
            chainid = element.start().get_parent().get_id()
            for res in element[2:-3]:
                stats.collect(element.type, res, chainid, self.structure.id)

def getid(fname):
    id = fname.split('.')
    if len(id) > 1:
        id = id[:-1]
    id = string.join(id, '.')
    return id

if __name__ == '__main__':
    pdblist = os.listdir(pdbdir)
    parser = PDBParser(PERMISSIVE=1)
    logger = LogFH(logfile, 'w')
    reader = DSSP_Reader()
    stats = StatCollector(taglist, vallist)
    for fname in pdblist:
        id = getid(fname)
        pdbfname = pdbdir+fname
        dsspfname = dsspdir+id+dssp_extension
        if not os.path.isfile(dsspfname): continue
        logger.log('processing '+id)
        structure = parser.get_structure(id, pdbfname)
        secstrdef = reader.extract_secstr(dsspfname)
        map = SSMap(structure, secstrdef, logger)
        map.run(stats)
    stats.show()
    logger.close()
