#!/usr/bin/env python
# F2mtz.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.
#
# 31st May 2006
#
# A wrapper for the CCP4 program f2mtz
#

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 F2mtz(DriverType = None):
    '''Create a F2mtz instance based on the passed in Driver type.'''

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

    class F2mtzWrapper(CCP4DriverInstance.__class__):
        '''A wrapper class for converting reflections to mtz.'''

        def __init__(self):
            CCP4DriverInstance.__class__.__init__(self)

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

            self._cell = None
            self._symmetry = None
            self._pname = None
            self._xname = None
            self._dname = None

        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_symmetry(self, symmetry):
            if type(symmetry) == type('str'):
                self._symmetry = symmetry.replace(' ', '')
            else:
                self._symmetry = str(symmetry)

        def xdsconv_anom2mtz(self):
            self.check_hklin()
            self.check_hklout()

            if self._symmetry is None:
                raise RuntimeError, 'symmetry not set'

            if self._cell is None:
                raise RuntimeError, 'cell not set'

            self.set_task('Converting reflections to mtz %s => %s' % \
                          (os.path.split(self.getHklin())[-1],
                           os.path.split(self.getHklout())[-1]))

            self.start()

            if self._pname and self._xname and self._dname:
                self.input('name project %s crystal %s dataset %s' % \
                           (self._pname, self._xname, self._dname))

            self.input('cell %f %f %f %f %f %f' % \
                       tuple(map(float, self._cell)))
            self.input('symmetry %s' % self._symmetry)
            self.input(
                'labout H K L IMEAN SIGIMEAN I(+) SIGI(+) I(-) SIGI(-)')
            self.input('CTYPOUT H H H J Q K M K M')

            self.close_wait()

            return self.get_ccp4_status()

        def xdsconv_nat2mtz(self):
            self.check_hklin()
            self.check_hklout()

            if self._symmetry is None:
                raise RuntimeError, 'symmetry not set'

            if self._cell is None:
                raise RuntimeError, 'cell not set'

            self.set_task('Converting reflections to mtz %s => %s' % \
                          (os.path.split(self.getHklin())[-1],
                           os.path.split(self.getHklout())[-1]))

            self.start()

            if self._pname and self._xname and self._dname:
                self.input('name prohect %s crystal %s dataset %s' % \
                           (self._pname, self._xname, self._dname))

            self.input('cell %f %f %f %f %f %f' % \
                       tuple(map(float, self._cell)))
            self.input('symmetry %s' % self._symmetry)
            self.input(
                'labout H K L IMEAN SIGIMEAN')
            self.input('CTYPOUT H H H J Q')

            self.close_wait()

            return self.get_ccp4_status()

        def f2mtz(self):
            self.check_hklin()
            self.check_hklout()

            if self._symmetry is None:
                raise RuntimeError, 'symmetry not set'

            if self._cell is None:
                raise RuntimeError, 'cell not set'

            self.set_task('Converting reflections to mtz %s => %s' % \
                          (os.path.split(self.getHklin())[-1],
                           os.path.split(self.getHklout())[-1]))

            self.start()

            if self._pname and self._xname and self._dname:
                self.input('name prohect %s crystal %s dataset %s' % \
                           (self._pname, self._xname, self._dname))

            self.input('cell %f %f %f %f %f %f' % \
                       tuple(map(float, self._cell)))
            self.input('symmetry %s' % self._symmetry)
            self.input('labout H K L FP FOM PHIB SIGFP')
            self.input('CTYPOUT H H H F W P Q')

            self.close_wait()

            return self.get_ccp4_status()


    return F2mtzWrapper()

if __name__ == '__main__':

    f = F2mtz()

    f.set_hklin('1vpj.phs')
    f.set_hklout('temp.mtz')
    f.set_cell([51.690000, 51.680000, 157.800000,
                90.000000, 90.000000, 90.000000])
    f.set_symmetry('P43212')

    status = f.f2mtz()

    print status