# $CVSHeader: pal8e/src/PSSD.py,v 1.40 2005/02/20 00:32:25 indraneel Rel $

import os
import sys
import string
import traceback
import ConfigParser
#import pypvm
import gzip
import bz2

from Bio.PDB.PDBParser import PDBParser
from Bio.PDB.Structure import Structure

from Exceptions import PSSDBuildException
import Parm
from Log import LogFH, PSSDLog
from Peptides import PSSDPeptides
from Definition import PSSDefinition
from Output import PSSDFileNames
from Output import PSSDPDBOut
from InsightScript import InsightScript
from PymolScript import PymolScript
from Sheet import PSSDSheets
from Plot import PSSDPlots
import Compare_DSSP
import Compare_stride
import Compare_sstruc
import Compare_pross
import Compare_molauto

class PSSD:
    '''
    Defines secondary structure for the given structure object.
    Call in this order:

    PSSD(conf) -- required, once only
    set_str(structure) -- for every new structure
    mkpssd([step]) -- for every new atructure
    '''

    ### Private methods

    def __init__(self, structure, header, trailer, conf, parm):
        '''
        Startup code.
        '''
        self.structure = structure
        self.header = header
        self.trailer = trailer
        self.conf = conf
        self.parm = parm
        
    def _find_chainstart(self, tag):
        for chain in self.structure[0].get_list():
            res = chain.get_list()[0]
            setattr(res, tag, 1)

    def _get_label(self):
        for line in self.header:
            if line[0:6] == 'HEADER':
                label = line[10:50]
                return label

    ### Public methods

    def run(self):
        '''
        Start the secondary structure definition on the supplied structure.
        '''
        self.logger = PSSDLog(self.conf)
        peptides = PSSDPeptides(self.structure, self.parm, self.logger)
        if not peptides:
            self.logger.close()
            return
        peptides.run_primary_ruleset()
        definition = PSSDefinition(self.parm, self.logger)
        definition.initiate_helix(peptides)
        ladders = peptides.build_quadruplets()
        ladders.initiate(peptides)
        definition.remove_hairpin_helix(peptides)
        ladders.extend(peptides, definition)
        peptides.mark_strands()
        definition.buildnew('E', peptides)
        definition.break_helices(peptides)
        definition.break_strands(peptides)
        definition.rollback_non_alpha_helices(peptides)
        sheets = PSSDSheets(self.logger)
        sheets.interpret(definition)
        
        peptides.delinvalidrulemarks({'H':['Hr1s','Hr2s','Hr1o','Hr2o'], \
                                      'E':['Er1s','Er2s','Er1o','Er2o']})

        definition.buildoverlap('H', 'E', 'HE')
        definition.buildoverlap('H', 'H', 'HH')
        definition.buildoverlap('E', 'E', 'EE')
        definition.buildstart('H', 'Hst', 1)
        definition.buildstart('E', 'Est', 1)
        self._find_chainstart('chainstart')

        taglist = ['HE','HH','EE','Hst','Est','Hr1s','Hr1o','Hr2s']
        taglist += ['Hr2o','Er1s','Er1o','Er2s','Er2o','blg','chainstart']
        taglist += ['q3']
        definition.buildnew(taglist, peptides)

        #peptides.initialise_tags(['compare'], dict)

	if (self.conf.has_option('Compare', 'use_dssp')
            and self.conf.getint('Compare', 'use_dssp')):
            compare_dssp = Compare_DSSP.DSSP(self.conf, self.logger)
            peptides.initialise_tags(['compare_dssp'])
            compare_dssp.compare(self.structure, peptides)
            definition.buildnew(['dssp_pssd_only_H', 'dssp_pssd_only_E', \
                                 'dssp_dssp_only'], peptides)

            definition.filter_length(['dssp_pssd_only_H', 'dssp_pssd_only_E', \
                                      'stride_pssd_only_H', 'stride_pssd_only_E'], 5)

	if (self.conf.has_option('Compare', 'use_stride')
            and self.conf.getint('Compare', 'use_stride')):
            compare_stride = Compare_stride.Stride(self.conf, self.logger)
            peptides.initialise_tags(['compare_stride'])
            compare_stride.compare(self.structure, peptides)
            #definition.buildnew(['stride_pssd_only_H', 'stride_pssd_only_E', \
            #                     'stride_stride_only'], peptides)

	if (self.conf.has_option('Compare', 'use_sstruc')
            and self.conf.getint('Compare', 'use_sstruc')):
            compare_sstruc = Compare_sstruc.Sstruc(self.conf, self.logger)
            peptides.initialise_tags(['compare_sstruc'])
            compare_sstruc.compare(self.structure, peptides)

	if (self.conf.has_option('Compare', 'use_pross')
            and self.conf.getint('Compare', 'use_pross')):
            compare_pross = Compare_pross.Pross(self.conf, self.logger)
            peptides.initialise_tags(['compare_pross'])
            compare_pross.compare(self.structure, peptides)

	if (self.conf.has_option('Compare', 'use_molauto')
            and self.conf.getint('Compare', 'use_molauto')):
            compare_molauto = Compare_molauto.Molauto(self.conf, self.logger)
            peptides.initialise_tags(['compare_molauto'])
            compare_molauto.compare(self.structure, peptides)

        definition.log()
        peptides.log()

        colors = [{'H':'0.0 1.0 1.0'}, \
                  {'E':'1.0 1.0 0'}, \
                  {'blg':'1.0 1.0 0.83'}, \
                  {'Hst':'0.0 1.0 0.0'}, \
                  {'Est':'0.0 1.0 0.0'}, \
                  {'HH':'0.83 1.0 0.83'}, \
                  {'EE':'0.83 1.0 0.83'}, \
                  {'HE':'1.0 0.0 1.0'}]

        plots = PSSDPlots(self.conf, self.logger)
        peptides.initialise_tags(['plot'], dict)
        peptides.initialise_tags(['pErr'])
        plots.set_colors(colors)
        plots.generate(sheets, definition)
        definition.buildnew(['pErr'], peptides)

        
        colors = [{'H':'cyan'}, \
                  {'E':'yellow'}, \
                  {'blg':'160,255,0'}, \
                  {'Hst':'green'}, \
                  {'Est':'green'}, \
                  {'HH':'0,255,127'}, \
                  {'EE':'0,255,127'}, \
                  {'HE':'magenta'}, \
                  {'pErr':'red'}]#, \
                  #{'q3':'blue'}]
        '''
        colors = [{'dssp_pssd_only_H':'cyan'}, \
                  {'dssp_pssd_only_E':'cyan'}, \
                  {'dssp_dssp_only':'red'}]
        
        colors = [{'stride_pssd_only_H':'cyan'}, \
                  {'stride_pssd_only_E':'cyan'}, \
                  {'stride_stride_only':'red'}]
        '''

        display = InsightScript(self.conf, self.logger)
        display.set_colors(colors)
        display.generate(definition)


        colors = [{'H':'cyan'}, \
                  {'E':'yellow'}, \
                  {'blg':'0.627,1.0,0.0'}, \
                  {'Hst':'green'}, \
                  {'Est':'green'}, \
                  {'HH':'0.0,1.0,0.5'}, \
                  {'EE':'0.0,1.0,0.5'}, \
                  {'HE':'magenta'}, \
                  {'pErr':'red'}]#, \
                  #{'q3':'blue'}]
        '''
        colors = [{'dssp_pssd_only_H':'cyan'}, \
                  {'dssp_pssd_only_E':'cyan'}, \
                  {'dssp_dssp_only':'red'}]

        colors = [{'stride_pssd_only_H':'cyan'}, \
                  {'stride_pssd_only_E':'cyan'}, \
                  {'stride_stride_only':'red'}]
        '''

        label = self._get_label()

        display = PymolScript(self.conf, self.logger)
        display.set_colors(colors)
        if label:
            display.set_label(label)
        display.generate(definition)

        definition.report_badstrands()

        #definition.stat()

        pdbout = PSSDPDBOut(self.structure, self.header,
                           self.trailer, self.conf)
        pdbout.write(definition)

        self.logger.close()
        return

class NonStopPSSD:
    '''
    Re-use same parser and config objects.
    '''
    def __init__(self):
        if not os.environ.has_key('PAL8ECONF'):
            os.environ['PAL8ECONF'] = os.path.expanduser('~/.pal8e.conf')
        conf_fname = os.getenv('PAL8ECONF')
        config = ConfigParser.ConfigParser()
        config.read(conf_fname)
        self.conf = config
        self.parser = PDBParser(PERMISSIVE=1)
        self.parm = Parm.PSSDParm(config)
        if config.has_option('Process', 'use_pvm') and \
               config.getint('Process', 'use_pvm') and \
               config.getint('Process', 'children'):
            self.using_pvm = 'true'
        else: self.using_pvm = None

    def _chk_compression(self, fname):
        if(fname[-2:] == '.Z'):
            return 'gzip'
        elif(fname[-3:] == '.gz'):
            return 'gzip'
        elif(fname[-4:] == '.bz2'):
            return 'bzip2'
        else: return None
        
    def _get_filehandle(self, fname, compression):
        # elif not reqd
        if compression == None:
            return file(fname)
        elif compression == 'gzip':
            return gzip.GzipFile(fname)
        elif compression == 'bzip2':
            return bz2.BZ2File(fname)

    ## Public methods

    def setnew(self, fname, structureID):
        self.structureID = structureID
        #print structureID
        self.conf.set('Input', 'fname', fname)
        self.conf.set('Input', 'structureID', structureID)

    def run(self):
        conf = self.conf
        fname = conf.get('Input', 'fname')
        if (conf.has_option('Output', 'overwrite')
            and (conf.getint('Output', 'overwrite')==0)):
            filenames = PSSDFileNames(conf)
            outfname = filenames.get_outfname()
            if os.path.exists(outfname): return None
        structureID = conf.get('Input', 'structureID')
        compression = self._chk_compression(fname)
        filehandle = self._get_filehandle(fname, compression)
        parser = self.parser
        try:
            structure = parser.get_structure(structureID, filehandle)
            header = parser.get_header()
            trailer = parser.get_trailer()
            pssd = PSSD(structure, header, trailer, conf, self.parm)
            pssd.run()
            return structure
        except:
            [etype, val, tb] = sys.exc_info()
            errors = traceback.format_exception(etype, val, tb)
            logfname = conf.get('Debug','dbgdir')+structureID
            compression = conf.getint('Output', 'compression')
            logger = LogFH(logfname, 'a', compression)
            logger.log(errors)
            logger.close()

    def run_using_pvm(self):
        if not self.using_pvm: return
        ptid = pypvm.parent()
        tid = pypvm.mytid()
        while 1:
            pypvm.recv(ptid, 1)
            args = pypvm.upkstr()
            args = string.split(args)
            if len(args) != 2: break
            [fname, structureID] = args
            self.setnew(fname, structureID)
            self.run()
            pypvm.initsend()
            pypvm.pkint([tid])
            pypvm.send(ptid, 1)
        pypvm.exit()
