#!/usr/bin/env python
# Cellparm.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 for the XDS jiffy program CELLPARM for computing average unit
# cell parameters, from the results of multiple CORRECT runs for input into
# a single XSCALE run.
#

import os
import sys
import math

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.Flags import Flags

def Cellparm(DriverType = None):

    DriverInstance = DriverFactory.Driver(DriverType)

    class CellparmWrapper(DriverInstance.__class__):
        '''A wrapper for wrapping CELLPARM.'''

        def __init__(self):

            # set up the object ancestors...

            DriverInstance.__class__.__init__(self)

            self.set_executable('cellparm')

            # input parameters
            self._cells = []
            self._n_refs = []

            return

        def add_cell(self, cell, n_ref):
            '''Add a unit cell which belongs to n_ref reflections.'''

            self._cells.append(cell)
            self._n_refs.append(n_ref)

            return

        def get_cell(self):
            '''Compute an average cell.'''

            if len(self._cells) < 1:
                raise RuntimeError, 'no input unit cell parameters'

            # check that the input cells are reasonably uniform -
            # be really relaxed and allow 5% variation!

            average_cell = [self._cells[0][j] for j in range(6)]
            number_cells = 1

            for j in range(1, len(self._cells)):
                cell = self._cells[j]
                for k in range(6):
                    average = average_cell[k] / number_cells
                    if math.fabs((cell[k] - average) / average) > 0.05 \
                           and not Flags.get_relax():
                        raise RuntimeError, 'incompatible unit cells'
                    average = average_cell[k] / number_cells
                    if math.fabs((cell[k] - average) / average) > 0.2:
                        raise RuntimeError, 'very incompatible unit cells'

                # it was ok to remember for later on..
                for k in range(6):
                    average_cell[k] += cell[k]
                number_cells += 1

            cellparm_inp = open(os.path.join(
                self.get_working_directory(), 'CELLPARM.INP'), 'w')

            for j in range(len(self._cells)):
                cell = self._cells[j]
                n_ref = self._n_refs[j]
                cellparm_inp.write('UNIT_CELL_CONSTANTS=')
                cellparm_inp.write(
                    '%.3f %.3f %.3f %.3f %.3f %.3f WEIGHT=%d\n' % \
                    (cell[0], cell[1], cell[2], cell[3], cell[4], cell[5],
                     n_ref))

            cellparm_inp.close()

            self.start()

            self.close_wait()

            # FIXME need to look for errors in here

            cellparm_lp = open(os.path.join(
                self.get_working_directory(), 'CELLPARM.LP'), 'r')
            data = cellparm_lp.readlines()

            return map(float, data[-1].split()[:6])

    return CellparmWrapper()

if __name__ == '__main__':
    cellparm = Cellparm()

    cellparm.add_cell(
        (64.585, 30.911, 34.886, 90.000, 105.822, 90.000), 284516)

    cellparm.add_cell(
        (64.584, 30.932, 34.879, 90.000, 105.832, 90.000), 32106)

    cellparm.add_cell(
        (64.658, 30.972, 34.918, 90.000, 105.823, 90.000), 1966)

    print cellparm.get_cell()