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

import os
import sys
from types import *
import ConfigParser
#import pypvm

from Bio.PDB.PDBParser import PDBParser

from PSSD import PSSD
from PSSD import NonStopPSSD
from Log import LogFH

class PSSDError(Exception):
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return self.value

class PSSDRunner:
    def __init__(self, config_fname_or_class):
        if type(config_fname_or_class) is StringType:
            self.conf_fname = conf_fname
        else: self.config = config_fname_or_class

    def _get_filelist(self, config):
        if hasattr(self, 'filelist'):
            return self.filelist
        if config.has_option('Input', 'pdbdir'):
            pdbdir = config.get('Input', 'pdbdir')
            config.remove_option('Input', 'pdbdir')
        else: pdbdir = ''

        lfd = None
        if config.has_option('Input', 'listfiledir'):
            lfd = config.get('Input', 'listfiledir')
            config.remove_option('Input', 'listfiledir')
        if config.has_option('Input', 'listfile'):
            lfn = config.get('Input','listfile')
            config.remove_option('Input', 'listfile')
        else:    
            lfn = os.getenv('HOSTNAME')
        lf = None
        if lfd:
            lfn = lfd + lfn
            lf = file(lfn)
        if lf:
            filelist = lf.readlines()
            lf.close()
            if pdbdir:
                for i in range(len(filelist)):
                    if filelist[i][:5] == 'break': break
                    filelist[i] = pdbdir+filelist[i][:-1]
        elif pdbdir:
            filelist = os.listdir(pdbdir)
            for i in range(len(filelist)):
                filelist[i] = pdbdir+filelist[i]
        else:
            return []
        return filelist

    def _get_filedict(self, config):
        filelist = self._get_filelist(config)

        if config.has_option('Input', 'filename_prefilter'):
            filename_prefilter = config.get('Input', 'filename_prefilter')
            prefil_len = len(filename_prefilter)
        else: filename_prefilter = None
        if config.has_option('Input', 'filename_postfilter'):
            filename_postfilter = config.get('Input', 'filename_postfilter')
            postfil_len = len(filename_postfilter)
        else: filename_postfilter = None

        filedict = {}
        for fname in filelist:
            if fname[:5] == 'break': break
            structureID = os.path.basename(fname)
            if filename_prefilter:
                if filename_prefilter != structureID[:prefil_len]:
                    continue
                structureID = structureID[prefil_len:]
            if filename_postfilter:
                if filename_postfilter != structureID[-postfil_len:]:
                    continue
                structureID = structureID[:-postfil_len]
            if len(structureID) == 0: continue
            filedict[structureID] = fname
        return filedict

    def _multiprocess_fork(self, config, fname, structureID):
        if len(self.children.keys()) >= self.multiprocess:
            pid,stat = os.wait()
            nonstop = self.children[pid]
            del self.children[pid]
        else:
            nonstop = NonStopPSSD()
            
        newpid = os.fork()
        if newpid == 0:
            nonstop.setnew(fname, structureID)
            nonstop.run()
            os._exit(0)
        else:
            self.logger.log(structureID+' '+fname+' (cid: '+str(newpid)+')')
            self.children[newpid] = nonstop

    def _multiprocess_pvm(self, config, fname, structureID):
        tid = -1
        if len(self.children.keys()) >= self.multiprocess:
            pypvm.recv()
            tid = pypvm.upkint(1)[0]
        else:
            args = ['PSSD.py']
            tid = pypvm.spawn('python', args, 0, '', 1)[0]
            #print 'spawned: ', tid
            if tid < 0:
                return
            self.children[tid] = 1
        self.logger.log(structureID+' '+fname+' (tid: '+str(tid)+')')
        pypvm.initsend()
        args = fname+' '+structureID
        pypvm.pkstr(args)
        pypvm.send(tid, 1)

    def run(self):
        if hasattr(self, 'conf_fname'):
            config = ConfigParser.ConfigParser()
            config.read(self.conf_fname)
            self.config = config

        config = self.config
        logfname = config.get('Debug','dbgdir')+config.get('Debug','logfile')
        compression = config.getint('Output', 'compression')
        self.logger = LogFH(logfname, compress=compression)
        config.remove_option('Debug', 'logfile')            

        filedict = self._get_filedict(config)
        if not filedict:
            self.logger.log('Empty filedict')
            return

        method = None
        if config.has_option('Process', 'children'):
            self.multiprocess = config.getint('Process', 'children')
        if self.multiprocess:
            self.children = {}
            if config.has_option('Process', 'use_pvm') and \
                   config.getint('Process', 'use_pvm'):
                method = self._multiprocess_pvm
            else:
                method = self._multiprocess_fork
        else:
            nonstop = NonStopPSSD()
        
        for structureID, fname in filedict.items():
            config.set('Input', 'structureID', structureID)
            config.set('Input', 'fname', fname)

            if method:
                method(config, fname, structureID)
            else:
                self.logger.log(structureID+' '+fname)
                nonstop.setnew(fname, structureID)
                nonstop.run()
                    
        if self.multiprocess:
            while len(self.children.keys()):
                if method == self._multiprocess_fork:
                    pid,stat = os.wait()
                    del self.children[pid]
                elif method == self._multiprocess_pvm:
                    pypvm.recv()
                    tid = pypvm.upkint(1)[0]
                    pypvm.initsend()
                    pypvm.pkstr(' ')
                    pypvm.send(tid, 1)
                    del self.children[tid]
        self.logger.close()

if __name__ == '__main__':
    if not os.environ.has_key('PAL8ECONF'):
        os.environ['PAL8ECONF'] = os.path.expanduser('~/.pal8e.conf')

    if len(sys.argv) == 1:
        conf_fname = os.getenv('PAL8ECONF')
        pssdrunner = PSSDRunner(conf_fname)
        pssdrunner.run()

    elif (2 <= len(sys.argv) <= 3):
        '''
        First argument is input filename with full path.
        Second argument if given is structure ID to be used.
        Structure ID if not given is first 4 characters of input filename.
        '''
        nonstop = NonStopPSSD()
        if nonstop.using_pvm:
            nonstop.run_using_pvm()
        else:
            fname = sys.argv[1]
            if len(sys.argv) == 3:
                structureID = sys.argv[2]
            else:
                structureID = os.path.basename(fname)[:4]
            nonstop.setnew(fname, structureID)
            nonstop.run()
        
    else:
        print >> sys.stderr, "python PAL8E.py <input PDB coordinate filename>"
        print >> sys.stderr, "or read the detailed instructions"
