#!/usr/bin/env python
# CCP4Decorator.py
#
#   Copyright (C) 2006 CCLRC, Graeme Winter
#
#   This code is distributed under the BSD license, a copy of which is 
#   included in the root directory of this package.
#
# 25th May 2006
# 
# A decorator to add hklin and hklout methods to a Driver instance.
# This will probably include some other interesting things like
# xyzin etc at some point in the future, once such things become 
# important.
# 
# Supported Keywords:
# HKLIN input MTZ reflection file (.mtz) 
# HLKOUT output MTZ reflection file (.mtz) 
# MAPIN input map file (.map) 
# MAPOUT output map file (.map) 
# XYZIN input coordinate file (.pdb) 
# XYZOUT output coordinate file (.pdb) 
# 
# All accessed via setHklin(hklin) getHklin() etc.
# 
# List from:
# 
# http://www.ccp4.ac.uk/dist/html/ccp4.html
# 
# Done:
# Add a mechanism for recording and parsing of loggraph output.
# 
# FIXME 06/NOV/06 would be fun to record the HKLIN and HKLOUT files to
#                 be able to track how the processing went and what jobs
#                 were executed to make that happen...
# 
# FIXME 24/NOV/06 need to be able to cope with columns in loggraph output
#                 running together.
# 

import os
import sys

from DecoratorHelper import inherits_from

def CCP4DecoratorFactory(DriverInstance):
    '''Create a CCP4 decorated Driver instance - based on the Driver
    instance which is passed in. This is an implementation of
    dynamic inheritance. Note well - this produces a new object and
    leaves the original unchanged.'''

    DriverInstanceClass = DriverInstance.__class__

    # verify that the input object satisfies the Driver interface -
    # in this case that at some point it inherited from there.
    # note well - instances of DefaultDriver must not be used
    # directly.

    if not inherits_from(DriverInstanceClass, 'DefaultDriver'):
        raise RuntimeError, 'object %s is not a Driver implementation' % \
              str(DriverInstance)

    # verify that the object matches the Driver specification

    class CCP4Decorator(DriverInstanceClass):
        '''A decorator class for a Driver object which will add some nice CCP4
        sugar on top.'''

        # to get hold of the original start() methods and so on - and
        # also to the parent class constructor
        _original_class = DriverInstanceClass

        def __init__(self):
            # note well - this is evil I am calling another classes constructor
            # in here. Further this doesn't know what it is constructing!
            
            self._original_class.__init__(self)
            
            self._hklin = None
            self._hklout = None
            self._xyzin = None
            self._xyzout = None
            self._mapin = None
            self._mapout = None

            # somewhere to store the loggraph output
            self._loggraph = { }
            
            return

        def set_hklin(self, hklin):
            return self.setHklin(hklin)

        def get_hklin(self):
            return self.getHklin()

        def check_hklin(self):
            return self.checkHklin()

        def setHklin(self, hklin):
            self._hklin = hklin
            return

        def getHklin(self):
            return self._hklin

        def checkHklin(self):
            if self._hklin is None:
                raise RuntimeError, 'hklin not defined'
            if not os.path.exists(self._hklin):
                raise RuntimeError, 'hklin %s does not exist' % self._hklin
        
        def set_hklout(self, hklout):
            return self.setHklout(hklout)

        def get_hklout(self):
            return self.getHklout()

        def check_hklout(self):
            return self.checkHklout()

        def setHklout(self, hklout):
            self._hklout = hklout
            return

        def getHklout(self):
            return self._hklout
        
        def checkHklout(self):
            if self._hklout is None:
                raise RuntimeError, 'hklout not defined'

            # check that these are different files!
            
            if self._hklout == self._hklin:
                raise RuntimeError, \
                      'hklout and hklin are the same file (%s)' % \
                      str(self._hklin)
        
        def set_xyzin(self, xyzin):
            return self.setXyzin(xyzin)

        def get_xyzin(self):
            return self.getXyzin()

        def check_xyzin(self):
            return self.checkXyzin()

        def setXyzin(self, xyzin):
            self._xyzin = xyzin
            return

        def getXyzin(self):
            return self._xyzin

        def checkXyzin(self):
            if self._xyzin is None:
                raise RuntimeError, 'xyain not defined'
            if not os.path.exists(self._xyzin):
                raise RuntimeError, 'xyzin %s does not exist' % self._xyzin

        def set_xyzout(self, xyzout):
            return self.setXyzout(xyzout)

        def get_xyzout(self):
            return self.getXyzout()

        def check_xyzout(self):
            return self.checkXyzout()
        
        def setXyzout(self, xyzout):
            self._xyzout = xyzout
            return

        def getXyzout(self):
            return self._xyzout
        
        def checkXyzout(self):
            if self._xyzout is None:
                raise RuntimeError, 'xyzout not defined'

        def set_mapin(self, mapin):
            return self.setMapin(mapin)

        def get_mapin(self):
            return self.getMapin()

        def check_mapin(self):
            return self.checkMapin()
        
        def setMapin(self, mapin):
            self._mapin = mapin
            return

        def getMapin(self):
            return self._mapin
        
        def checkMapin(self):
            if self._mapin is None:
                raise RuntimeError, 'mapin not defined'
            if not os.path.exists(self._mapin):
                raise RuntimeError, 'mapin %s does not exist' % self._mapin
        
        def set_mapout(self, mapout):
            return self.setMapout(mapout)

        def get_mapout(self):
            return self.getMapout()

        def check_mapout(self):
            return self.checkMapout()

        def setMapout(self, mapout):
            self._mapout = mapout
            return

        def getMapout(self):
            return self._mapout

        def checkMapout(self):
            if self._mapout is None:
                raise RuntimeError, 'mapout not defined'
        
        def describe(self):
            '''An overloading of the Driver describe() method.'''

            description = 'CCP4 program: %s' % self.get_executable()

            if not self._hklin is None:
                description += ' %s' % ('hklin')
                description += ' %s' % (self._hklin)

            if not self._hklout is None:
                description += ' %s' % ('hklout')
                description += ' %s' % (self._hklout)

            if not self._xyzin is None:
                description += ' %s' % ('xyzin')
                description += ' %s' % (self._xyzin)

            if not self._xyzout is None:
                description += ' %s' % ('xyzout')
                description += ' %s' % (self._xyzout)

            if not self._mapin is None:
                description += ' %s' % ('mapin')
                description += ' %s' % (self._mapin)

            if not self._mapout is None:
                description += ' %s' % ('mapout')
                description += ' %s' % (self._mapout)

            return description
        
        def start(self):
            '''Add all the hklin etc to the command line then call the
            base classes start() method. Also make any standard ccp4
            scratch directories. The latter shouldnt be needed however.'''

            for env in ['BINSORT_SCR', 'CCP4_SCR']:
                if os.environ.has_key(env):
                    directory = os.environ[env]                    
                    self.add_scratch_directory(directory)
                    try:
                        os.mkdir(directory)
                    except:
                        pass

            if not self._hklin is None:
                self.add_command_line('hklin')
                self.add_command_line(self._hklin)

            if not self._hklout is None:
                self.add_command_line('hklout')
                self.add_command_line(self._hklout)

            if not self._xyzin is None:
                self.add_command_line('xyzin')
                self.add_command_line(self._xyzin)

            if not self._xyzout is None:
                self.add_command_line('xyzout')
                self.add_command_line(self._xyzout)

            if not self._mapin is None:
                self.add_command_line('mapin')
                self.add_command_line(self._mapin)

            if not self._mapout is None:
                self.add_command_line('mapout')
                self.add_command_line(self._mapout)

            # delegate the actual starting to the parent class
            self._original_class.start(self)

            return

        def check_ccp4_errors(self):
            '''Look through the standard output for a few "usual" CCP4
            errors, for instance incorrect file formats &c.'''

            # check that the program has finished...

            if not self.finished():
                raise RuntimeError, 'program has not finished'

            for line in self.get_all_output():
                if 'CCP4 library signal' in line:
                    error = line.split(':')[1].strip()
                    
                    # throw away the "status" in brackets

                    if '(' in error:
                        error = error.split('(')[0]

                    # handle specific cases...

                    if 'Write failed' in error:
                        # work out why...
                        for l in self.get_all_output():
                            if '>>>>>> System signal' in l:
                                cause = l.split(':')[1].split('(')[0]
                                raise RuntimeError, '%s:%s' % (error, cause)

                    # then cope with the general case
                    
                    else:
                        raise RuntimeError, error

            return 

        def get_ccp4_status(self):
            '''Check through the standard output and get the program
            status. Note well - this will only work if called once the
            program is finished.'''

            # check that the program has finished...

            if not self.finished():
                raise RuntimeError, 'program has not finished'

            # look in the last 10 lines for the status

            # Added 1/SEP/06 to check if .exe is on the end of the
            # command line...

            program_name = os.path.split(self.get_executable())[-1].lower()
            if program_name[-4:] == '.exe':
                program_name = program_name[:-4]

            # special case for FFT which calls itself FFTBIG in the status
            # output..

            if program_name == 'fft':
                program_name = 'fftbig'

            for line in self.get_all_output()[-10:]:
                l = line.split()
                if len(l) > 1:
                    if l[0][:-1].lower() == program_name:
                        # then this is the status line
                        status = line.split(':')[1].replace('*', '')
                        return status.strip()

                    if l[0][:-1].lower() == program_name.split('-')[0]:
                        # then this is also probably the status line
                        status = line.split(':')[1].replace('*', '')
                        return status.strip()

            raise RuntimeError, 'could not find status'

        def parse_ccp4_loggraph(self):
            '''Look through the standard output of the program for
            CCP4 loggraph text. When this is found store it in a
            local dictionary to allow exploration.'''

            # reset the loggraph store
            self._loggraph = { }

            output = self.get_all_output()

            for i in range(len(output)):
                line = output[i]
                if '$TABLE' in line:

                    n_dollar = line.count('$$')
                    
                    current = line.split(':')[1].replace('>',
                                                                  '').strip()
                    self._loggraph[current] = { }
                    self._loggraph[current]['columns'] = []
                    self._loggraph[current]['data'] = []

                    loggraph_info = ''

                    # FIXME 09/FEB/07 this is assumed to represent the
                    # end of a loggraph but really it isn't

                    # almost exactly one year later this has come back to
                    # bite me - with chef! indeed... OK, in Smartie the
                    # way this worked was to count the number of $$
                    # following the $TABLE spell - when this was equal to
                    # 4 this table is complete. Good idea.

                    # while not 'inline graphs' in line and not 'FONT' in line:
                    while n_dollar < 4:
                        n_dollar += line.count('$$')
                        loggraph_info += line

                        if n_dollar == 4:
                            break
                        
                        i += 1
                        line = output[i]

                    # at this stage I should have the whole 9 yards in
                    # a single string... 

                    tokens = loggraph_info.split('$$')
                    self._loggraph[current]['columns'] = tokens[1].split()

                    if len(tokens) < 4:
                        raise RuntimeError, 'loggraph "%s" broken' % current
                    
                    data = tokens[3].split('\n')

                    # pop takes the data off the end so...
                    # data.reverse()

                    columns = len(self._loggraph[current]['columns'])

                    # while len(data) > 0:
                    # record = []
                    # for i in range(columns):
                    # record.append(data.pop())
                    # self._loggraph[current]['data'].append(record)

                    # code around cases where columns merge together...

                    for j in range(len(data)):
                        record = data[j].split()
                        if len(record) == columns:
                            self._loggraph[current]['data'].append(record)

            return self._loggraph                
        
    return CCP4Decorator()

if __name__ == '__main__':
    if not os.environ.has_key('XIA2CORE_ROOT'):
        raise RuntimeError, 'XIA2CORE_ROOT not defined'

    sys.path.append(os.path.join(os.environ['XIA2CORE_ROOT'],
                                 'Python',
                                 'Driver'))
    
    from DriverFactory import DriverFactory    

    d = DriverFactory.Driver('script')

    d = CCP4DecoratorFactory(d)

    from pydoc import help

    print help(d.__class__)
    
