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

import os
from Bio.PDB.PDBIO import PDBIO
import gzip

pdbio = PDBIO()

class PSSDFileNames:
    def __init__(self, conf):
        self.conf = conf

    def get_outfname(self):
        conf = self.conf
        secoutdir = conf.get('Output','secoutdir')
        structureID = conf.get('Input', 'structureID')
        endtag = conf.get('Output', 'secfilename_posttag')
        outfname = secoutdir+structureID+endtag
        return outfname

class PSSDPDBOut(PSSDFileNames):
    def __init__(self, structure, header, trailer, conf):
        self.conf = conf
        PSSDFileNames.__init__(self, conf)
        self.outfname = self.get_outfname()
        if os.path.isfile(self.outfname):
            os.remove(self.outfname)        
        self.structure = structure
        self.header = header
        self.trailer = trailer

    def _get_definition_lines(self, definition):
        lines = []
        helix_ctr = 0
        definition.sort_by_position()
        for type in ['H', 'E', 'S']:
            for elem in definition:
                if elem.type != type: continue
                if elem.type == 'H':
                    helix_ctr += 1
                    lines += [elem.get_H_line(helix_ctr)]
                elif elem.type == 'E':
                    lines += [elem.get_E_line()]
                elif elem.type == 'S':
                    lines += [elem.get_S_line()]
        return lines

    def _rescmp(self, resa, resb):
        return cmp(resa.strpos, resb.strpos)

    def _get_pair_lines(self, pairs):
        lines = []
        pair_ctr = 0
        keylist = pairs.keys()
        keylist.sort(self._rescmp)
        for res in keylist:
            for pair in pairs[res]:
                pair_ctr += 1
                line = "PAIRS "+" "+"   "+" "+"   "+" "
                line += res.get_resname() + " "
                line += res.get_parent().get_id() + " "
                line += "%4d" % (res.get_id()[1]) + res.get_id()[2] + " "
                line += pair.get_resname() + " "
                line += pair.get_parent().get_id() + " "
                line += "%4d" % pair.get_id()[1]
                line += pair.get_id()[2]
                line += '\n'
                lines += [line]
        return lines

    def _write(self, lines, header, trailer):
        sidx = 0
        for (sidx,line) in enumerate(header):
            if line[0:6] == "HELIX " or \
               line[0:6] == "SHEET ":
                break
        if sidx:
            eidx = 0
            for line in header:
                if line[0:6] == "HELIX " or \
                       line[0:6] == "SHEET ":
                    eidx = header.index(line)
        if sidx:
            header = header[:sidx] + header[eidx+1:]
        else:
            for (sidx,line) in enumerate(header):
                if line[0:6] == "TURN  " or \
                       line[0:6] == "SSBOND" or \
                       line[0:6] == "LINK  ":
                    break
        if sidx:
            header = header[:sidx] + lines + header[sidx:]
        else:
            header = header + lines
        pdbio.set_structure(self.structure)
        body = pdbio.get_lines()
        compression = self.conf.getint('Output', 'compression')
        if compression:
            outfh = gzip.GzipFile(self.outfname+'.gz', 'w')
        else:
            outfh = file(self.outfname, 'w')
        outfh.writelines(header+body+trailer)
        outfh.close()

    # Public Methods

    def write(self, definition):
        lines = self._get_definition_lines(definition)
        pairs = definition.get_E_pairs()
        lines += self._get_pair_lines(pairs)
        self._write(lines, self.header, self.trailer)

