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

import os
import re
import string
import gzip

class PymolScript:
    def __init__(self, conf, logger):
        self.conf = conf
        self.logger = logger
        self.pmlfilereqd = conf.getboolean('Output', 'pmlfilereqd')
        self.objID = 0
        self.label = ''
        structureID = conf.get('Input', 'structureID')
        self.outfname = conf.get('Output', 'pmloutdir')+structureID+'.pml'
        if self.pmlfilereqd 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', 'pmlpdbdir')+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 _pmlstart(self):
        structureID = self.conf.get('Input', 'structureID')
        self.id = structureID+'_'+str(self.objID)
        lines = ['load '+self.relfname+', '+self.id+'\n']
        lines += ['hide everything, '+self.id+'\n']
        lines += ['show ribbon, '+self.id+'\n']
        lines += ['set ribbon_sampling, 1, '+self.id+'\n']
        lines += ['set ribbon_width, 2, '+self.id+'\n']
        lines += ['set orthoscopic, 1\n']
        lines += ['zoom '+self.id+', complete=1\n']
        lines += ['color gray, (name CA in '+self.id+')\n']
        lines += ['origin '+self.id+'\n']
        lines += ['set sphere_scale, 0.15\n']
        lines += ['show spheres, (name CA in '+self.id+')\n']
        if self.label:
            lines += ['set_title '+self.id+', 1, '+self.label[:31]+'\n']
        self.outbuf = lines

    def _colorbytype(self, type, color, definition):
        if re.compile(',').search(color):
            self.outbuf += ['set_color custom_color, ['+color+']\n']
            color = 'custom_color'
        for elem in definition:
            if elem.type != type: continue
            chain = elem.start().get_parent().id
            if chain == ' ': chain = "''"
            stres = elem.start().id[1]
            endres = elem.end().id[1]
            line = 'color '+color+', (chain '+chain+' and resi '
            line += str(stres)+'-'+str(endres)+' and name CA) in '+self.id+'\n'
            self.outbuf += [line]

    def _showpairs(self, definition):
        self.outbuf += ['set dash_length, 0.5\n']
        self.outbuf += ['set dash_gap, 0.1\n']
        self.outbuf += ['set dash_width, 1.0\n']
        d = 0
        for elem in definition:
            if elem.type != 'E': continue
            for res in elem:
                types = [[res.pairs['q1'], 'red']]
                types += [[res.pairs['q2'], 'orange']]
                types += [[res.pairs['q3'], 'blue']]
                for (pairs, color) 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 = 'dist d'+str(d)+', '
                        line += '(chain '+stch+' and resi '+str(res.id[1])
                        line += ' and name CA in '+self.id+'), '
                        line += '(chain '+ench+' and resi '+str(pair.id[1])
                        line += ' and name CA in '+self.id+')\n'
                        lines = [line]
                        lines += ['hide label, d'+str(d)+'\n']
                        lines += ['color '+color+', d'+str(d)+'\n']
                        d += 1
                        self.outbuf += lines

    def _pmlend(self, definition):
        lines = []
        for elem in definition:
            if elem.type == 'chainstart':
                ch = elem.start().get_parent().id
                if ch == ' ': ch = "''"
                resid = elem.start().id[1]
                selection = '(chain '+ch+' and resi '+str(resid)
                selection += ' and name CA) in '+self.id
                lines += ['label '+selection+', (chain,resi)\n']
        self.outbuf += lines
        self.outbuf += ['#stereo crosseye\n']
        self.outbuf += ['#stereo swap\n']

    ## Public Methods

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

    def set_label(self, label):
        self.label = label

    def generate(self, definition):
        if not self.pmlfilereqd: return
        self._pmlstart()
        for item in self.color:
            for (type, color) in item.items():
                self._colorbytype(type, color, definition)
        self._showpairs(definition)
        self._pmlend(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
