#!/usr/bin/env python
# XDSConv.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.
#
# A wrapper to run xdsconv
#

import os
import sys
import shutil

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'

if not os.path.join(os.environ['XIA2CORE_ROOT'],
                    'Python') in sys.path:
    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 Handlers.Streams import Debug

def XDSConv(DriverType = None):

    DriverInstance = DriverFactory.Driver(DriverType)

    class XDSConvWrapper(DriverInstance.__class__):
        '''A wrapper for wrapping XDSCONV.'''

        def __init__(self):
            DriverInstance.__class__.__init__(self)
            self.set_executable('xdsconv')

            self._input_file = None
            self._cell = None
            self._symmetry = None
            self._output_file = None

            return

        def set_input_file(self, input_file):
            self._input_file = input_file
            return

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

        def get_cell(self):
            return self._cell

        def get_symmetry(self):
            return self._symmetry

        def set_symmetry(self, symmetry):
            self._symmetry = symmetry
            return

        def set_output_file(self, output_file):
            self._output_file = output_file
            return

        def parse_xds_ascii(self, file):
            '''Parse the XDS ascii file for interesting things.'''

            results = { }

            for line in open(file, 'r').readlines():
                if not line[0] == '!':
                    break

                if '!FORMAT' in line:
                    tokens = line.split()
                    for t in tokens:
                        if 'MERGED' in t:
                            if t.split('=')[-1].lower() == 'false':
                                raise RuntimeError, 'input unmerged'
                        if 'FRIEDEL' in t:
                            results['friedel'] = t.split('=')[-1].lower()

                if '!UNIT_CELL_CONSTANTS' in line:
                    results['cell'] = tuple(map(float,
                                                line.split()[1:]))

                if '!INCLUDE_RESOLUTION' in line:
                    results['resolution_range'] = map(float,
                                                      line.split()[1:])

                if '!SPACE_GROUP' in line:
                    results['spacegroup'] = int(line.split()[-1])

            return results

        def convert(self):
            if not self._input_file:
                raise RuntimeError, 'no input file specified'

            if not self._output_file:
                raise RuntimeError, 'no output file specified'

            # make the output file link a relative rather than
            # absolute path... FIXME this is unix specific!
            if self.get_working_directory() in self._output_file:
                self._output_file = self._output_file.replace(
                    self.get_working_directory(), './')

            # perhaps move input file to CWD

            if len(self._input_file) > 49:
                if len(os.path.split(self._input_file)[-1]) > 49:
                    raise RuntimeError, 'input file name too long'

                shutil.copyfile(
                    self._input_file,
                    os.path.join(self.get_working_directory(),
                                 os.path.split(self._input_file)[-1]))

                self._input_file = os.path.split(self._input_file)[-1]

            header = self.parse_xds_ascii(self._input_file)

            if not self._cell:
                self._cell = header['cell']

            if not self._symmetry:
                self._symmetry = header['spacegroup']

            if header.has_key('resolution_range'):
                self._resolution = header['resolution_range']
            else:
                self._resolution = [100.0, 0.1]
            self._resolution.sort()
            self._resolution.reverse()

            inp = open(os.path.join(
                self.get_working_directory(), 'XDSCONV.INP'), 'w')

            inp.write('SPACE_GROUP_NUMBER=%d\n' % self._symmetry)

            inp.write(
                'UNIT_CELL_CONSTANTS= %.2f %.2f %.2f %.2f %.2f %.2f\n' % \
                self._cell)

            inp.write('INPUT_FILE=%s XDS_ASCII %.2f %.2f\n' % \
                      (self._input_file, self._resolution[0],
                       self._resolution[1]))

            inp.write('OUTPUT_FILE=%s IALL FRIEDEL\'S_LAW=%s\n' % \
                      (self._output_file, header['friedel'].upper()))

            inp.close()

            self.start()
            self.close_wait()

            # should really parse the output

            return

    return XDSConvWrapper()

# need to add a test here...