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

# Generates Insight log files colored according to PSSD definition
# PSSD output files must be available
#

import os
import string

from Bio.PDB.PDBParser import PDBParser
from Bio.PDB.Polypeptide import CaPPBuilder

from Comparisons import PSSD

class Insight:
    def getElements(self, fname):
        pssd = PSSD()
        pssd.store(fname)
        return pssd.oldelements

    def store(self, fname):
  	ppb = CaPPBuilder(radius=5.0)
        parser = PDBParser(PERMISSIVE=1)
        structure = parser.get_structure('xx',fname)
        self.structure = structure
	peptides = ppb.build_peptides(structure)
        self.peptides = peptides

    def getStructure(self):
        return self.structure
    
    def getPeptides(self):
        return self.peptides

    def getStartLines(self, relfname, id, position):
        lines = ['Get Molecule PDB User '+relfname+' '+id+' Heteroatom\n']
        lines += ['Display Molecule Only Atoms Trace '+id+'\n']
        lines += ['Color Molecule Atoms '+id+ \
                  ' Specified Specification 247,247,247\n']
        lines += ['Display Molecule On Atoms All '+id+':*H\n']
        lines += ['Display Molecule Off Atoms All '+id+':HOH\n']
        lines += ['Color Molecule Atoms '+id+':*H'+ \
                  ' Specified Specification 255,0,255\n']
        pos = str(position * -150)
        lines += ['Move '+id+' Center_of_mass Relative 0,'+pos+ \
                  ',0 Screen\n']
        lines += ['Clip -Auto 400 0 -Perspective INSIGHTII\n']
        return lines

    def getElementLines(self, elements, id):
        lines = []
        for element in elements:
            if (element.type == 'H'):
                color = '0,255,255'
            elif (element.type == 'E'):
                color = '255,255,0'
            else: continue
            stch = endch = element.chain
            if stch == ' ': stch = endch = ''
            stres = element.start
            endres = element.end
            sticode = element.starticode
            endicode = element.endicode
            if sticode == ' ': sticode = ''
            if endicode == ' ': endicode = ''
            line = 'Color Molecule Atoms '+id+':'
            line += stch + str(stres) + sticode + '-'
            line += endch + str(endres) + endicode
            line += ' Specified Specification '+color+'\n'
            lines += [line]
        return lines

    def getJoinLines(self, peptides, id):
        '''bond fragments'''
        color = '0,255,0'
        lines = []
        for (i,peptide) in enumerate(peptides[:-1]):
            stch = peptide[-1].get_parent().id
            endch = peptides[i+1][0].get_parent().id
            if (stch != endch): continue
            if stch == ' ': stch = ''
            if endch == ' ': endch = ''
            starticode = peptide[-1].id[2]
            endicode = peptides[i+1][0].id[2]
            if starticode == ' ': starticode = ''
            if endicode == ' ': endicode = ''
            stid = str(peptide[-1].id[1])
            endid = str(peptides[i+1][0].id[1])
            line = 'Bond Create -Fragment_Window Single '
            line += id+':'+stch+stid+starticode+':CA '
            line += id+':'+endch+endid+endicode+':CA '
            line += '-Assimilate\n'
            line += 'Color Molecule Atoms '+id+':'
            line += stch+stid+starticode+'-'
            line += endch+endid+endicode+':CA '
            line += 'Specified Specification '+color+'\n'

            lines += [line]
        return lines

    def getFragmentN(self, peptides, id):
        '''color first CA of fragment'''
        color = '0,255,0'
        lines = []
        for peptide in peptides:
            ch = peptide[0].get_parent().id
            if ch == ' ': ch = ''
            resicode = peptide[0].id[2]
            if resicode == ' ': resicode = ''
            rid = peptide[0].id[1]
            line = 'Color Molecule Atoms '+id+':'
            line += ch + str(rid) + resicode + ':CA'
            line += ' Specified Specification '+color+'\n'
            lines += [line]
        return lines

    def getFirstElementLines(self, elements, id):
        '''color first atom of first element'''
        color = '255,0,0'
        lines = []
        element = elements[0]
        stch = element.chain
        if stch == ' ': stch = ''
        stres = element.start
        sticode = element.starticode
        if sticode == ' ': sticode = ''
        line = 'Color Molecule Atoms '+id+':'
        line += stch + str(stres) + sticode
        line += ' Specified Specification '+color+'\n'
        lines += [line]
        return lines

    def getChainN(self, id):
        '''color first CA of every chain'''
        color = '255,0,0'
        lines = []
        structure = self.getStructure()
        model = structure.get_list()[0]
        for chain in model.get_list():
            chid = chain.id
            if chid == ' ': chid = ''
            res = chain.get_list()[0]
            (resha, resid, resic) = res.id
            if resic == ' ': resic = ''
            line = 'Color Molecule Atoms '+id+':'
            line += chid + str(resid) + resic + ':CA'
            line += ' Specified Specification '+color+'\n'
            lines += [line]
        return lines
        
    def getInsLogLines(self, pssdFname, relfname, id, position):
        elements = self.getElements(pssdFname)
        self.store(pssdFname)
        peptides = self.getPeptides()
        lines = []
        lines += self.getStartLines(relfname, id, position)
        lines += self.getElementLines(elements, id)
        lines += self.getJoinLines(peptides, id)
        #lines += self.getFragmentN(peptides, id)
        #lines += self.getFirstElementLines(elements, id)
        lines += self.getChainN(id)
        lines += ['\n']
        return lines

    def mkInsLog(self, pssdFname, outFname, relfname, id, position=0):
        lines = self.getInsLogLines(pssdFname, relfname, id, position)
        fh = file(outFname, 'a')
        fh.writelines(lines)
        fh.close()
        
class Interface:
    def __init__(self):
        self.position = -1
        self.counter = -1
        self.fnamectr = 0
    
    def runOneFile(self, pssdFname, outFname, relfname):
        base = os.path.basename(pssdFname)
        base = base.split('.')
        base = base[0]
        id = 'P_' + base
        self.position += 1
        insight = Insight()
        insight.mkInsLog(pssdFname, outFname, relfname, id, self.position)

    def runMany(self, listfile, pssdDir, outDir, relDir):
        #files = os.listdir(pssdDir)
        lfh = file(listfile)
        files = lfh.readlines()
        lfh.close()
        for (i, fn) in enumerate(files):
            files[i] = fn[:-1]
        for fname in files:
            pssdFname = pssdDir+fname
            relFname = relDir+fname
            self.counter += 1
            if self.counter == 9:
                self.fnamectr += 1
                self.counter = -1
                self.position = -1
            outFname = outDir + "%04d" % self.fnamectr + '.log'
            self.runOneFile(pssdFname, outFname, relFname)

if __name__ == '__main__':
    interface = Interface()
    listfile = './files'
    pssdDir = '/mnt/e12/data/indraneel/scopPSSD+HET/'
    outDir = '/mnt/e12/data/indraneel/scopPSSD+HETlog/'
    relDir = 'scopPSSD+HET/'
    interface.runMany(listfile, pssdDir, outDir, relDir)
