#!/usr/bin/env python
# Mtz2various.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.
#
# 27th October 2006
#
# A wrapper for the CCP4 program mtz2various, for converting mtz files
# to .sca format.
#

import os
import sys

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'))

from Driver.DriverFactory import DriverFactory
from Decorators.DecoratorFactory import DecoratorFactory

def Mtz2various(DriverType = None):
    '''A factory for Mtz2variousWrapper classes.'''

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

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

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

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

            # this will allow extraction of specific intensities
            # from a multi-set reflection file
            self._dataset_suffix = ''

            return

        def set_suffix(self, suffix):
            if suffix:
                self._dataset_suffix = '_%s' % suffix
            else:
                self._dataset_suffix = suffix
            return

        def convert(self):
            '''Convert the input reflection file to .sca format.'''

            self.check_hklin()
            self.check_hklout()

            self.start()

            labin = 'I(+)=I(+)%s SIGI(+)=SIGI(+)%s ' % \
                    (self._dataset_suffix, self._dataset_suffix)
            labin += 'I(-)=I(-)%s SIGI(-)=SIGI(-)%s' % \
                     (self._dataset_suffix, self._dataset_suffix)

            self.input('output scal')
            self.input('labin %s' % labin)

            self.close_wait()

            output = self.get_all_output()

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

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

        def convert_shelx(self, unmerged = False):
            '''Convert the input reflection file to SHELX hklf4 format.'''

            self.check_hklin()
            self.check_hklout()

            self.start()

            if self._dataset_suffix or unmerged:
                labin = 'I=I%s SIGI=SIGI%s' % \
                        (self._dataset_suffix, self._dataset_suffix)

            else:
                labin = 'I=IMEAN SIGI=SIGIMEAN'

            self.input('output shelx')
            self.input('labin %s' % labin)

            self.close_wait()

            output = self.get_all_output()

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

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

    return Mtz2variousWrapper()

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

    for hklin in sys.argv[1:]:
        m2v = Mtz2various()
        m2v.set_hklin(hklin)
        m2v.set_hklout('%s.sca' % hklin[:-4])
        m2v.convert()