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

import os
import time
import gzip

from Element import PSSDElements

class PSSDPlot(PSSDElements):
    def __init__(self, conf, logger):
        self.conf = conf
        self.logger = logger
        PSSDElements.__init__(self, logger)

    def _get_X_coord(self, elem, pair):
        hasx = None
        for res in elem:
            if res.plot.has_key(self.sheetID):
                hasx = res.plot[self.sheetID][0]
                break
        if not hasx:
            neighbours = []
            neighbours += elem.start().prev[1:3]
            neighbours += elem.start().next[1:3]
            neighbours += elem.end().prev[1:3]
            neighbours += elem.end().next[1:3]
            for neighbour in neighbours:
                if neighbour == None: continue
                if neighbour.plot.has_key(self.sheetID):
                    hasx = neighbour.plot[self.sheetID][0]
                    break
        space = [1,1]
        x = None
        for res in pair:
            if not elem in res.elempairs: continue
            [x,y] = res.plot[self.sheetID]
            for p in res.pairs['all']:
                if p in elem: continue
                if p.plot.has_key(self.sheetID) and \
                       p.plot[self.sheetID]:
                    if p.plot[self.sheetID][0] == (x-1):
                        space[0] = 0
                    if p.plot[self.sheetID][0] == (x+1):
                        space[1] = 0
                    if space == [0,0]:
                        m = 'No space for:'
                        self.logger.log('trace', m, 'ploterr')
                        elem.log('ploterr')
                        elem.start().pErr = 1
                        self.Error = 'No space for '+str(elem)
                        self.Error += ' beside '+str(pair)
                        return
        if x == None:
            m = 'elem not in res.elempairs for res in pair: check: '
            elem.log('ploterr')
            self.logger.log('trace', m, 'ploterr')
            elem.start().pErr = 1
            self.Error = str(elem)+' not in res.elempairs for res in '
            self.Error += str(pair)
            return
        if (hasx == (x-1)) and (space[0] == 1):
            return hasx
        elif (hasx == (x+1)) and (space[1] == 1):
            return hasx
        if space[0] == 1:
            x -= 1
        else: x += 1
        return x

    def _set_res_coord(self, res, coords):
        if res.plot[self.sheetID]:
            m = 'placing '
            m += res.get_parent().id+str(res.id)
            m += ' again on sheet '+self.sheetID
            self.logger.log('trace', m, 'ploterr')
        if (coords[0], coords[1]) in self.plottedcoords:
            p = self.plottedcoords[(coords[0], coords[1])]
            m = 'Sheet: '+self.sheetID+' '
            m += res.get_parent().id+str(res.id)
            m += ' overwriting '+p.get_parent().id+str(p.id)
            self.logger.log('trace', m, 'ploterr')
            res.pErr = 1
            p.pErr = 1
        res.plot[self.sheetID] = coords
        self.plottedcoords[(coords[0], coords[1])] = res

    def _mkcoords(self, elem):
        for pair in elem.pairs:
            if pair.plotted: break
        else:
            # first strand on sheet
            x,y = 0,0.0
            for res in elem:
                res.plot[self.sheetID] = [x,y]
                y += 1.0
            elem.plotted = 1
            return
        # check for empty space on either side of pair
        x = self._get_X_coord(elem, pair)
        if self.Error: return
        for res in elem:
            if not res.plot.has_key(self.sheetID):
                res.plot[self.sheetID] = None
        # plot paired residues
        for res in pair:
            for p in res.pairs['all']:
                if p in elem:
                    y = res.plot[self.sheetID][1]
                    if p.plot[self.sheetID] == [x,y]: continue
                    self._set_res_coord(p, [x,y])
        # find first plotted residue in this column
        for i in range(len(elem)):
            if elem[i].plot[self.sheetID] and \
                   elem[i].plot[self.sheetID][0] == x:
                break
        # fill in bulges
        for res in elem[i+1:-1]:
            if res.plot[self.sheetID]: continue
            if res.prev[1].plot[self.sheetID] and \
               res.next[1].plot[self.sheetID] and \
               res.prev[1].plot[self.sheetID][0] == x and \
               res.next[1].plot[self.sheetID][0] == x:
                y = (res.prev[1].plot[self.sheetID][1] + \
                     res.next[1].plot[self.sheetID][1]) / 2.0
                self._set_res_coord(res, [x,y])
        # fill in edges
        for j in range(len(elem)-1, -1, -1):
            if elem[j].plot[self.sheetID] and \
                   elem[j].plot[self.sheetID][0] == x:
                break
        iy = elem[i].plot[self.sheetID][1]
        jy = elem[j].plot[self.sheetID][1]
        if (iy < jy): shift = 1.0
        else: shift = -1.0
        if not (i == 0):
            for k in range(i-1, -1, -1):
                res = elem[k]
                y = res.next[1].plot[self.sheetID][1]
                self._set_res_coord(res, [x, y-shift])
        if not (j == len(elem)-1):
            for k in range(j+1, len(elem)):
                res = elem[k]
                y = res.prev[1].plot[self.sheetID][1]
                self._set_res_coord(res, [x, y+shift])
        X = elem.start().plot[self.sheetID][0]
        m = ''
        for res in elem[1:]:
            if res.plot[self.sheetID]:
                x = res.plot[self.sheetID][0]
                if x != X:
                    m = 'non-uniform X coord in element: '
            else:
                self.Error = 'incomplete X coord for '+str(elem)
                m = 'missing X coord in element (multiple res bulge): '
            if m:
                self.logger.log('trace', m, 'ploterr')
                elem.log('ploterr')
                res.pErr = 1
                if self.Error: return
        elem.plotted = 1

    def _process(self, elem):
        self._mkcoords(elem)
        if self.Error: return
        for pair in elem.pairs:
            if pair.plotted: continue
            if pair.sheetID != self.sheetID: continue
            self._process(pair)
            if self.Error: return

    def _shift_coords(self, pos):
        coords = {}
        for elem in self:
            for res in elem:
                if res.plot.has_key(self.sheetID) and \
                       res.plot[self.sheetID]:
                    coords[res.plot[self.sheetID][pos]] = 1
        klist = coords.keys()
        klist.sort()
        for (i, k) in zip(range(len(klist)), klist):
            coords[k] = i
        if pos == 0:
            self.bbox[pos] = ((i+1) * 40) + 40
        elif pos == 1:
            self.bbox[pos] = ((i+1) * 20) + 40
        for elem in self:
            for res in elem:
                if res.plot.has_key(self.sheetID) and \
                       res.plot[self.sheetID]:
                    k = res.plot[self.sheetID][pos]
                    if coords.has_key(k):
                        p = coords[k]
                        if pos == 0:
                            c = 40 + (p * 40)
                        elif pos == 1:
                            c = 40 + (p * 20)
                        res.plot[self.sheetID][pos] = c

    def _boxres(self, res, color, boxcolor):
        [x,y] = res.plot[self.sheetID]
        label = ''
        cid = res.get_parent().id
        if cid != ' ':
            label += '0.83 0.17 0.0 setrgbcolor\n(' + cid + ') show\n'
        label += '0 setgray (' + str(res.id[1]) + ') show\n'
        self.outbuf += [color + ' setrgbcolor\n']
        self.outbuf += ["newpath "+str(x)+' '+str(y)+" moveto fillbox "]
        if boxcolor:
            self.outbuf += [boxcolor + ' setrgbcolor\n']
            self.outbuf += ["newpath "+str(x)+' '+str(y)+" moveto box "]
        self.outbuf += ["newpath "+str(x+1)+' '+str(y+2)+" moveto\n"]
        self.outbuf += [label]

    def _boxup(self, rescolors):
        self.outbuf += ["(Times-Roman) findfont 10 scalefont setfont\n"]
        for elem in self:
            for res in elem:
                if res.plot.has_key(self.sheetID) and \
                       res.plot[self.sheetID]:
                    if res.pErr:
                        boxcolor = '1.0 0.0 0.0'
                    else: boxcolor = None
                    self._boxres(res, rescolors[res], boxcolor)

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

    def _linkres(self, resi, resj):
        [x1, y1] = resi.plot[self.sheetID]
        [x2, y2] = resj.plot[self.sheetID]
        if abs(x1 - x2) > 40:
            self.outbuf += ['(Times-Roman) findfont 7 scalefont setfont\n']
            self.outbuf += ["newpath "+str(x1+15)+' '+str(y1)+" moveto "]
            self.outbuf += [' ('+str(resj.id[1])+') show\n']
            self.outbuf += ['(Times-Roman) findfont 10 scalefont setfont\n']
            return
        if self._rescmp(resi, resj) > -1: return
        if x1 < x2: x1 += 30
        elif x2 < x1: x2 += 30
        y1 += 7
        y2 += 7
	self.outbuf += ["newpath "+str(x1)+' '+str(y1)+" moveto "]
	self.outbuf += [str(x2)+' '+str(y2)+" lineto closepath stroke\n"]

    def _link_consecutive(self, resi, resj):
        [x1, y1] = resi.plot[self.sheetID]
        [x2, y2] = resj.plot[self.sheetID]
        if y1 < y2: y1 += 15
        elif y2 < y1: y2 += 15
        else: return
        if abs(x1 - x2) > 40: return
        x1 += 15
        x2 += 15
	self.outbuf += ["newpath "+str(x1)+' '+str(y1)+" moveto "]
	self.outbuf += [str(x2)+' '+str(y2)+" lineto closepath stroke\n"]

    def _linkup(self):
        deflst = [['q1', '1.0 0.0 0.0'], \
                  ['q2', '0.0 0.0 1.0'], \
                  ['q3', '0.0 1.0 1.0']]
        for (pairtype, color) in deflst:
            self.outbuf += [color+' setrgbcolor\n']
            for elem in self:
                for res in elem:
                    if not res.plot.has_key(self.sheetID): continue
                    if not res.plot[self.sheetID]: continue
                    for pair in res.pairs[pairtype]:
                        if pair.plot.has_key(self.sheetID) and \
                               pair.plot[self.sheetID]:
                            self._linkres(res, pair)
        self.outbuf += ['0.0 1.0 0.0 setrgbcolor\n']
        for elem in self:
            for res in elem[:-1]:
                if res.plot.has_key(self.sheetID) and \
                       res.plot[self.sheetID] and \
                       res.next[1].plot.has_key(self.sheetID) and \
                       res.next[1].plot[self.sheetID]:
                    self._link_consecutive(res, res.next[1])

    def _error_message(self):
        m = 'Error plotting elements:'
        self.outbuf += ["(Times-Roman) findfont 12 scalefont setfont\n"]
        self.outbuf += ['newpath 10 740 moveto 0 setgray ('+m+') show\n']
        self.outbuf += ["(Times-Roman) findfont 10 scalefont setfont\n"]
        y = 725
        for elem in self:
            m = 'newpath '+str(10)+' '+str(y)+' moveto 0 setgray ('
            m += str(elem)
            m += ') show\n'
            self.outbuf += [m]
            y -= 12
        if self.Error:
            m = 'newpath '+str(10)+' '+str(y)+' moveto 0 setgray ('
            m += self.Error
            m += ') show\n'
            self.outbuf += [m]            
                
    ## Public Methods

    def generate(self, sheet):
        '''
        A single residue can occur at multiple places, on this plot as well
        as any other.
        This calculates only the position(s) of each residue.
        '''
        for element in sheet:
            self.append(element)
        self.sheetID = sheet.sheetID
        self.sort_by_length()
        self.reverse()
        self.plottedcoords = {}
        for element in sheet:
            element.plotted = None
        self.Error = None
        self._process(self[0])
        self.bbox = [None, None]
        self._shift_coords(0)
        self._shift_coords(1)
        return self.Error

    def getepsheader(self):
        bbox = '30 30 '+str(self.bbox[0])+' '+str(self.bbox[1])
        structureID = self.conf.get('Input', 'structureID')
        lines = []
        lines += ["%!PS-Adobe-2.0 EPSF-2.0\n"]
        lines += ["%%Title:"+structureID+'_'+self.sheetID+"\n"]
        lines += ["%%Creator: pssd\n"]
        lines += ["%%CreationDate: "+ time.ctime() +"\n"]
        lines += ["%%DocumentFonts: Times-Roman\n"]
        lines += ["%%BoundingBox: "+bbox+' \n']
        lines += ["%%EndComments\n"]
        lines += ["/box { 30 0 rlineto 0 15 rlineto -30 0 rlineto\n"]
        lines += ["\tclosepath stroke } def\n"]
        lines += ["/fillbox { 30 0 rlineto 0 15 rlineto -30 0 rlineto\n"]
        lines += ["\tclosepath fill stroke } def\n"]
        lines += ["/up { 4 10 rlineto 4 -10 rlineto closepath\n"]
        lines += ["\t0.5 setgray fill stroke } def\n"]
        lines += ["/down { 4 -10 rlineto 4 10 rlineto closepath\n"]
        lines += ["\t0.5 setgray fill stroke } def\n"]
        lines += ["1.0 setlinewidth\n"]
        return lines

    def getpsheader(self, ctr):
        lines = []
        structureID = self.conf.get('Input', 'structureID')
        label = structureID+'_'+self.sheetID
        page = str(ctr)
        lines += ["%%Page: "+page+' '+page+"\n"]
        lines += ["(Times-Roman) findfont 20 scalefont setfont\n"]
        lines += ["newpath 10 760 moveto (" +label+ ")show\n"]
        return lines

    def getpstrailer(self):
        return ["showpage\n"]

    def getlines(self, rescolors):
        '''
        This links each residue to its pair at the closest location.
        For pairs that exist at multiple locations on this sheet
        write the attached residueID beside the non-closest locations.
        Mark colors appropriately for each residue and link.
        '''
        self.outbuf = []
        if self.Error: self._error_message()
        self._boxup(rescolors)
        self._linkup()
        return self.outbuf

class PSSDPlots(list):
    def __init__(self, conf, logger):
        self.conf = conf
        self.logger = logger
        self.psfilereqd = conf.getboolean('Output', 'psfilereqd')
        structureID = conf.get('Input', 'structureID')
        psoutdir = conf.get('Output', 'psoutdir')
        self.outfID = psoutdir+structureID

    def _getpsheader(self):
        structureID = self.conf.get('Input', 'structureID')
        lines = []
        lines += ["%!PS-Adobe-2.0\n"]
        lines += ["%%Title:"+structureID+"\n"]
        lines += ["%%Creator: pssd\n"]
        lines += ["%%CreationDate: "+ time.ctime() +"\n"]
        lines += ["%%DocumentFonts: Times-Roman\n"]
        lines += ["%%Orientation: Portrait\n"]
        lines += ["%%DocumentMedia: Letter 600 800\n"]
        lines += ["%%Pages: (atend)\n"]
        lines += ["%%EndComments\n"]
        lines += ["/box { 30 0 rlineto 0 15 rlineto -30 0 rlineto\n"]
        lines += ["\tclosepath stroke } def\n"]
        lines += ["/fillbox { 30 0 rlineto 0 15 rlineto -30 0 rlineto\n"]
        lines += ["\tclosepath fill stroke } def\n"]
        lines += ["/up { 4 10 rlineto 4 -10 rlineto closepath\n"]
        lines += ["\t0.5 setgray fill stroke } def\n"]
        lines += ["/down { 4 -10 rlineto 4 10 rlineto closepath\n"]
        lines += ["\t0.5 setgray fill stroke } def\n"]
        lines += ["1.0 setlinewidth\n"]
        return lines

    def _getpstrailer(self):
        lines = []
        lines += ["%%Pages: " + str(len(self)) + '\n']
        return lines

    def _get_final_colors(self, colors, definition):
        rescolors = {}
        for item in colors:
            for (type, color) in item.items():
                for elem in definition:
                    if elem.type != type: continue
                    for res in elem:
                        rescolors[res] = color
        return rescolors

    # Public Methods

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

    def generate(self, sheets, definition):
        if not self.psfilereqd: return
        rescolors = self._get_final_colors(self.color, definition)
        compression = self.conf.getint('Output', 'compression')
        psfname = self.outfID+'.ps'
        if compression:
            outfh = gzip.GzipFile(psfname+'.gz', 'w')
        else:
            outfh = file(psfname, 'w')
        outfh.writelines(self._getpsheader())
        ctr = 0
        for sheet in sheets:
            ctr += 1
            plot = PSSDPlot(self.conf, self.logger)
            plot.generate(sheet)
            self.append(plot)
            plotlines = plot.getlines(rescolors)
            pslines = plot.getpsheader(ctr)
            pslines += plotlines
            pslines += plot.getpstrailer()
            outfh.writelines(pslines)
            epslines = plot.getepsheader()
            epslines += plotlines
            epsfname = self.outfID+'_'+plot.sheetID+'.eps'
            if compression:
                epsfh = gzip.GzipFile(epsfname+'.gz', 'w')
            else:
                epsfh = file(epsfname, 'w')
            epsfh.writelines(epslines)
            epsfh.close()
        outfh.writelines(self._getpstrailer())
        outfh.close()
