#!/usr/bin/env python
# Combat.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.
#
# 5th June 2006
#
# An example of an combat CCP4 program wrapper, which can be used as the
# base for other wrappers.
#
# Provides:
#
# Conversion from XDS format and Unmerged scalepack format to MTZ
# unmerged format.
#

import os
import sys

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

if not os.environ.has_key('XIA2_ROOT'):
    raise RuntimeError, 'XIA2_ROOT not defined'

sys.path.append(os.path.join(os.environ['XIA2CORE_ROOT'],
                             'Python'))

if not os.environ['XIA2_ROOT'] in sys.path:
    sys.path.append(os.environ['XIA2_ROOT'])

from Driver.DriverFactory import DriverFactory
from Decorators.DecoratorFactory import DecoratorFactory
from Handlers.Streams import Debug

def Combat(DriverType = None):
    '''A factory for CombatWrapper classes.'''

    DriverInstance = DriverFactory.Driver(DriverType)
    CCP4DriverInstance = DecoratorFactory.Decorate(DriverInstance, 'ccp4')

    class CombatWrapper(CCP4DriverInstance.__class__):
        '''A wrapper for Combat, using the CCP4-ified Driver.'''

        def __init__(self):
            # generic things
            CCP4DriverInstance.__class__.__init__(self)

            self.set_executable(os.path.join(
                os.environ.get('CBIN', ''), 'combat'))

            self._pname = None
            self._xname = None
            self._dname = None

            # optional information - useful if this is
            # a scalepack input file though as we need to get the
            # spacegroup and so on set...
            self._spacegroup = None
            self._cell = None
            self._wavelength = None

            return

        def _find_largest_negative_intensity_xds(self):
            '''Find the largest negative intensity in the reflection
            file - this will be used to define a scale factor.'''

            imin = 0.0

            for record in open(self.get_hklin(), 'r').readlines():
                if record[0] == '!':
                    continue
                i = float(record.split()[3])
                if i < imin:
                    imin = i

            return imin

        def set_project_info(self, pname, xname, dname):
            self._pname = pname
            self._xname = xname
            self._dname = dname
            return

        def set_cell(self, cell):
            self._cell = cell

        def set_spacegroup(self, spacegroup):
            self._spacegroup = spacegroup

        def set_wavelength(self, wavelength):
            self._wavelength = wavelength

        def run(self):
            '''Actually convert to MTZ.'''

            self.check_hklin()
            self.check_hklout()

            # inspect hklin to decide what format it is...

            format = None

            # check for XDS format

            first_char = open(self.get_hklin(), 'r').read(1)
            if first_char == '!':
                # may be XDS
                first_line = open(self.get_hklin(), 'r').readline()
                if 'XDS_ASCII' in first_line:
                    format = 'XDSASCII'

            # check for scalepack format...?

            if not format:
                format = 'SCAL_NM2'

            if not format:
                raise RuntimeError, 'unknown format input file %s' % \
                      self.get_hklin()

            if format == 'SCAL_NM2' and not self._cell:
                raise RuntimeError, 'need CELL for scalepack unmerged'
            if format == 'SCAL_NM2' and not self._spacegroup:
                raise RuntimeError, 'need SPACEGROUP for scalepack unmerged'

            self.start()
            self.input('input %s' % format)
            if format == 'XDSASCII':

                # determine the appropriate scale factor
                imin = self._find_largest_negative_intensity_xds()

                Debug.write('Imin found to be %f' % imin)

                scale = 1.0
                while imin * scale < -999999.0:
                    scale *= 0.1

                Debug.write('Determined scale factor of %f' % scale)
                self.input('scale %s' % scale)

            if self._pname and self._xname and self._dname:
                self.input('pname %s' % self._pname)
                self.input('xname %s' % self._xname)
                self.input('dname %s' % self._dname)

            if self._cell:
                self.input('cell %f %f %f %f %f %f' % self._cell)
            if self._spacegroup:
                self.input('symmetry %s' % self._spacegroup)
            if self._wavelength:
                self.input('wavelength %f' % self._wavelength)

            self.close_wait()

            # check the status

            try:
                self.check_for_errors()
                self.check_ccp4_errors()

            except RuntimeError, e:
                try:
                    os.remove(self.get_hklout())
                except:
                    pass

                raise e

            return

    return CombatWrapper()

if __name__ == '__main__':
    # run a test

    # test XDS_ASCII
    try:
        if len(sys.argv) > 1:
            hklin = sys.argv[1]
        else:
            hklin = 'XDS_ASCII.HKL'

        c = Combat()
        c.write_log_file('combat-debug.log')
        c.set_hklin(hklin)
        c.set_hklout('temp.mtz')
        c.run()
    except RuntimeError, e:
        print e

    if False:

        # test unmerged polish

        c = Combat()

        hklin = os.path.join(os.environ['X2TD_ROOT'],
                             'Test', 'UnitTest', 'Interfaces',
                             'Scaler', 'Unmerged', 'TS00_13185_unmerged_INFL.sca')

        c.set_hklin(hklin)
        c.set_project_information('TS00', '13185', 'INFL')
        c.set_spacegroup('P212121')
        c.set_cell((57.74, 73.93, 86.57, 90.00, 90.00, 90.00))
        c.set_hklout('TS00_13185_unmerged_INFL.mtz')
        c.run()