#!/usr/bin/env python
# Cad.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.
#
# 26th October 2006
#
# A wrapper for the CCP4 program Cad, for merging multiple reflection files
# into a single reflection file, e.g. in preparation for phasing.
#
# This will take a single HKLOUT file and a list of HKLIN files, c/f the
# Sortmtz wrapper implementation.
#
# Ok, this Cad wrapper will handle two cases. The first is relabeling all
# columns by e.g. appending a new suffix to the name (this should default
# to the dname I guess) and setting the unit cell to some sanctioned value.
#
# The second aspect to this is the merging of the reflection files into
# on big one.
#
# This should almost be implemented as two wrappers - oh well! Aha - implement
# two methods - merge & update - this can handle the two distinct use cases.

import os
import sys

if not os.environ.has_key('XIA2CORE_ROOT'):
    raise RuntimeError, 'XIA2CORE_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'))

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

# locally required wrappers

from Mtzdump import Mtzdump

# external functionality

from Modules.FindFreeFlag import FindFreeFlag

def Cad(DriverType = None):
    '''A factory for CadWrapper classes.'''

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

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

        def __init__(self):
            # generic things
            CCP4DriverInstance.__class__.__init__(self)
            self.set_executable(os.path.join(
                os.environ.get('CBIN', ''), 'cad'))

            self._hklin_files = []

            self._new_cell_parameters = None
            self._new_column_suffix = None

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

            # stuff to specifically copy in the freer column...
            self._freein = None
            self._freein_column = 'FreeR_flag'

            return

        def add_hklin(self, hklin):
            '''Add a reflection file to the list to be sorted together.'''
            self._hklin_files.append(hklin)
            return

        def set_freein(self, freein):

            # I guess I should check in here that this file actually
            # exists... - also that it has a sensible FreeR column...

            if not os.path.exists(freein):
                raise RuntimeError, 'reflection file does not exist: %s' % \
                      freein

            cname = FindFreeFlag(freein)

            Debug.write('FreeR_flag column identified as %s' % cname)

            self._freein = freein
            self._freein_column = cname

            return

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

        def set_new_suffix(self, suffix):
            '''Set a column suffix for this dataset.'''
            self._new_column_suffix = suffix
            return

        def set_new_cell(self, cell):
            '''Set a new unit cell for this dataset.'''
            self._new_cell_parameters = cell
            return

        def merge(self):
            '''Merge multiple reflection files into one file.'''

            if not self._hklin_files:
                raise RuntimeError, 'no hklin files defined'

            self.check_hklout()

            hklin_counter = 0

            # for each reflection file, need to gather the column names
            # and so on, to put in the cad input here - also check to see
            # if the column names clash... check also that the spacegroups
            # match up...

            spacegroup = None
            column_names = []
            column_names_by_file = { }

            for hklin in self._hklin_files:
                md = Mtzdump()
                md.set_working_directory(self.get_working_directory())
                md.set_hklin(hklin)
                md.dump()
                columns = md.get_columns()
                spag = md.get_spacegroup()

                if spacegroup is None:
                    spacegroup = spag

                if spag != spacegroup:
                    raise RuntimeError, 'spacegroups do not match'

                column_names_by_file[hklin] = []

                for c in columns:
                    name = c[0]
                    if name in ['H', 'K', 'L']:
                        continue
                    if name in column_names:
                        raise RuntimeError, 'duplicate column names'
                    column_names.append(name)
                    column_names_by_file[hklin].append(name)

            # if we get to here then this is a good set up...

            # create the command line

            hklin_counter = 0
            for hklin in self._hklin_files:
                hklin_counter += 1
                self.add_command_line('hklin%d' % hklin_counter)
                self.add_command_line(hklin)

            self.start()

            hklin_counter = 0

            for hklin in self._hklin_files:
                column_counter = 0
                hklin_counter += 1
                labin_command = 'labin file_number %d' % hklin_counter
                for column in column_names_by_file[hklin]:
                    column_counter += 1
                    labin_command += ' E%d=%s' % (column_counter, column)

                self.input(labin_command)

            self.close_wait()

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

            except RuntimeError, e:
                # something went wrong; remove the output file
                try:
                    os.remove(self.get_hklout())
                except:
                    pass
                raise e

            return self.get_ccp4_status()

        def update(self):
            '''Update the information for one reflection file.'''

            if not self._hklin_files:
                raise RuntimeError, 'no hklin files defined'

            if len(self._hklin_files) > 1:
                raise RuntimeError, 'can have only one hklin to update'

            hklin = self._hklin_files[0]

            self.check_hklout()

            column_names_by_file = { }
            dataset_names_by_file = { }

            md = Mtzdump()
            md.set_hklin(hklin)
            md.dump()
            columns = md.get_columns()

            column_names_by_file[hklin] = []
            dataset_names_by_file[hklin] = md.get_datasets()

            # get a dataset ID - see FIXME 03/NOV/06 below...

            dataset_ids = [md.get_dataset_info(d)['id'] for \
                           d in md.get_datasets()]

            for c in columns:
                name = c[0]
                if name in ['H', 'K', 'L']:
                    continue

                column_names_by_file[hklin].append(name)

            self.add_command_line('hklin1')
            self.add_command_line(hklin)
            self.start()

            dataset_id = dataset_ids[0]

            if self._pname and self._xname and self._dname:
                self.input('drename file_number 1 %d %s %s' % \
                           (dataset_id, self._xname, self._dname))
                self.input('dpname file_number 1 %d %s' % \
                           (dataset_id, self._pname))

            column_counter = 0
            labin_command = 'labin file_number 1'
            for column in column_names_by_file[hklin]:
                column_counter += 1
                labin_command += ' E%d=%s' % (column_counter, column)

            self.input(labin_command)

            # FIXME perhaps - ASSERT that we want only the information from
            # the first dataset here...

            pname, xname, dname = dataset_names_by_file[hklin][0].split('/')
            dataset_id = dataset_ids[0]

            # FIXME 03/NOV/06 this needs to id the dataset by it's number
            # not by pname/xname/dname, as the latter get's confused if the
            # xname is a number...

            if self._new_cell_parameters:
                a, b, c, alpha, beta, gamma = self._new_cell_parameters
                self.input('dcell file_number 1 %d %f %f %f %f %f %f' % \
                           (dataset_id, a, b, c, alpha, beta, gamma))

            if self._new_column_suffix:
                suffix = self._new_column_suffix
                column_counter = 0
                labout_command = 'labout file_number 1'
                for column in column_names_by_file[hklin]:
                    column_counter += 1
                    labout_command += ' E%d=%s_%s' % \
                                     (column_counter, column, suffix)

                self.input(labout_command)

            self.close_wait()

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

            except RuntimeError, e:
                # something went wrong; remove the output file
                try:
                    os.remove(self.get_hklout())
                except:
                    pass
                raise e

            return self.get_ccp4_status()

        def copyfree(self):
            '''Copy the free column from freein into hklin -> hklout.'''

            if not self._hklin_files:
                raise RuntimeError, 'no hklin files defined'

            if len(self._hklin_files) > 1:
                raise RuntimeError, 'can have only one hklin to update'

            hklin = self._hklin_files[0]

            # get the resolution limit to give as a limit for the FreeR
            # column

            md = Mtzdump()
            md.set_working_directory(self.get_working_directory())
            md.set_hklin(hklin)
            md.dump()
            resolution_range = md.get_resolution_range()
            
            self.check_hklout()
            if self._freein is None:
                raise RuntimeError, 'freein not defined'
            if self._freein_column is None:
                raise RuntimeError, 'freein column not defined'

            self.add_command_line('hklin1')
            self.add_command_line(self._freein)
            self.add_command_line('hklin2')
            self.add_command_line(hklin)
            self.start()

            self.input('labin file_number 1 E1=%s' % self._freein_column)
            self.input('resolution file_number 1 %f %f' % resolution_range)
            self.input('labin file_number 2 all')

            self.close_wait()

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

            except RuntimeError, e:
                # something went wrong; remove the output file
                try:
                    os.remove(self.get_hklout())
                except:
                    pass
                raise e

            return self.get_ccp4_status()

    return CadWrapper()

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

    c = Cad()
    c.add_hklin(os.path.join(os.environ['X2TD_ROOT'],
                             'Test', 'UnitTest', 'Wrappers', 'Cad',
                             'TS01_12847_truncated.mtz'))
    c.set_new_suffix('NATIVE')
    average_unit_cell = (228.21, 52.61, 44.11, 90.00, 100.64, 90.00)
    c.set_new_cell(average_unit_cell)
    c.set_hklout('bar.mtz')
    c.update()