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

import os
import string
import gzip

class InsightScript:
    def __init__(self, conf, logger):
        self.conf = conf
        self.logger = logger
        self.insfilereqd = conf.getboolean('Output', 'insfilereqd')
        self.objID = 0
        structureID = conf.get('Input', 'structureID')
        self.outfname = conf.get('Output', 'insoutdir')+structureID+'.log'
        if self.insfilereqd and os.path.isfile(self.outfname):
            os.remove(self.outfname)
        self._set_relfname(conf)

    def _set_relfname(self, conf):
        fname = os.path.basename(conf.get('Input', 'fname'))
        relfname = conf.get('Display', 'inspdbdir')+fname
        relfname = relfname.split('.')
        tag = relfname[-1]
        if (tag == 'Z') or \
           (tag == 'gz') or \
           (tag == 'bz2'):
            relfname = relfname[:-1]
        relfname = string.join(relfname, '.')
        self.relfname = relfname

    def _insstart(self):
        structureID = self.conf.get('Input', 'structureID')
        self.id = 'P_'+structureID+'_'+str(self.objID)
        lines = ['Get Molecule PDB User '+self.relfname+' '+self.id+'\n']
        lines += ['Display Molecule Only Atoms Trace '+self.id+'\n']
        lines += ['Color Molecule Atoms '+self.id+ \
                  ' Specified Specification 247,247,247\n']
        pos = str(self.objID * -150)
        lines += ['Move '+self.id+' Center_of_mass Relative 0,'+pos+ \
                  ',0 Screen\n']
        lines += ['Clip -Auto 400 0 -Perspective INSIGHTII\n']
        self.outbuf = lines

    def _colorbytype(self, type, color, definition):
        for elem in definition:
            if elem.type != type: continue
            stch = elem.start().get_parent().id
            endch = elem.end().get_parent().id
            stres = elem.start().id[1]
            endres = elem.end().id[1]
            if stch == ' ': stch = ''
            if endch == ' ': endch = ''
            line = 'Color Molecule Atoms '+self.id+':'
            line += stch + str(stres) + '-' + endch + str(endres)
            line += ' Specified Specification '+color+'\n'
            self.outbuf += [line]

    def _showpairs(self, definition):
        for elem in definition:
            if elem.type != 'E': continue
            for res in elem:
                types = [[res.pairs['q1'], 'Triple']]
                types += [[res.pairs['q2'], 'Double']]
                types += [[res.pairs['q3'], 'Partial Double']]
                for (pairs, btype) in types:
                    for pair in pairs:
                        stch = res.get_parent().id
                        ench = pair.get_parent().id
                        if stch > ench: continue
                        if (stch == ench) and \
                           (res.id[1] > pair.id[1]): continue
                        if stch == ' ': stch = ''
                        if ench == ' ': ench = ''
                        line = '#Bond Create "'+btype+'"'
                        line += self.id+':'+stch+str(res.id[1])+':CA '
                        line += self.id+':'+ench+str(pair.id[1])+':CA'
                        line += '\n'
                        self.outbuf += [line]

    def _insend(self, definition):
        for elem in definition:
            if elem.type == 'chainstart':
                ch = elem.start().get_parent().id
                if ch == ' ': ch = ''
                resid = elem.start().id[1]
                line = 'Render Molecule CPK '+self.id+':'
                line += ch+str(resid)+':CA 0.5 Medium\n'
                self.outbuf += [line]
        self.outbuf += ['#Stereo Stereo_Mode 3.5 Side_By_Side 0.25\n']

    ## Public Methods

    def set_colors(self, color):
        self.color = color

    def generate(self, definition):
        if not self.insfilereqd: return
        self._insstart()
        for item in self.color:
            for (type, color) in item.items():
                self._colorbytype(type, color, definition)
        self._showpairs(definition)
        self._insend(definition)
        compression = self.conf.getint('Output', 'compression')
        if compression:
            outfh = gzip.GzipFile(self.outfname+'.gz', 'a')
        else:
            outfh = file(self.outfname, 'a')
        outfh.writelines(self.outbuf)
        outfh.close()
        self.objID += 1
