#!/usr/bin/env python
# Rebatch.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.
#
# 21/SEP/06
#
# A wrapper for the CCP4 program REBATCH.
#
# FIXME 30/NOV/06 need to add a facility to use this to select batches
#                 from the input reflection file...
#

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

def Rebatch(DriverType = None):
    '''A factory for RebatchWrapper classes.'''

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

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

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

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

            self._first_batch = 0
            self._add_batch = 0

            self._exclude_batches = []

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

            return

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

        def set_first_batch(self, first_batch):
            self._first_batch = first_batch
            return

        def set_add_batch(self, add_batch):
            self._add_batch = add_batch
            return

        def exclude_batch(self, batch):
            if not batch in self._exclude_batches:
                self._exclude_batches.append(batch)
            return

        def exclude_batches(self):

            self.check_hklin()
            self.check_hklout()

            if not self._exclude_batches:
                raise RuntimeError, 'no batches to exclude'

            self.start()

            for batch in sorted(self._exclude_batches):
                self.input('batch %d reject' % batch)

            self.close_wait()

            # check for errors...
            try:
                self.check_for_errors()
                self.check_ccp4_errors()

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

                raise e

            output = self.get_all_output()

            return

        def limit_batches(self, first, last):
            '''Limit the batches to first to last.'''

            self.check_hklin()
            self.check_hklout()

            if self._first_batch > 0 or self._add_batch > 0:
                raise RuntimeError, 'limiting when batches set to renumber'

            self.start()

            if first - 1 > 0:
                self.input('batch 0 to %d reject' % (first - 1))
            self.input('batch %d to 1000000 reject' % (last + 1))

            self.close_wait()

            # check for errors...
            try:
                self.check_for_errors()
                self.check_ccp4_errors()

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

                raise e

            # get out the new batch range...

            output = self.get_all_output()

            min = 10000000
            max = -10000000

            for i in range(len(output)):
                o = output[i]
                if o.split()[:5] == ['Old', 'batch', 'New', 'batch', 'Max']:
                    j = i + 2
                    m = output[j]
                    while not 'SUMMARY_END' in m:
                        l = m.split()
                        if len(l) == 3:
                            batch = int(l[1])
                            if batch < min:
                                min = batch
                            if batch > max:
                                max = batch
                        j += 1
                        m = output[j]

            new_batches = (min, max)

            return new_batches

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

            if self._first_batch > 0 and self._add_batch > 0:
                raise RuntimeError, 'both first and add specified'

            if self._first_batch == 0 and self._add_batch == 0:
                raise RuntimeError, 'neither first nor add specified'

            self.start()

            if self._first_batch > 0:
                command = 'batch all start %d' % self._first_batch
            else:
                command = 'batch all add %d' % self._add_batch

            self.input(command)

            if self._pname and self._xname and self._dname:
                command = 'batch all pname %s xname %s dname %s' % \
                          (self._pname, self._xname, self._dname)
                self.input(command)

            self.close_wait()

            # check for errors...
            try:
                self.check_for_errors()
                self.check_ccp4_errors()

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

                raise e

            # get out the new batch range...

            output = self.get_all_output()

            min = 10000000
            max = -10000000

            for i in range(len(output)):
                o = output[i]
                if o.split()[:5] == ['Old', 'batch', 'New', 'batch', 'Max']:
                    j = i + 2
                    m = output[j]
                    while not 'SUMMARY_END' in m:
                        l = m.split()
                        if len(l) >= 3:
                            batch = int(l[1])
                            if batch < min:
                                min = batch
                            if batch > max:
                                max = batch
                        j += 1
                        m = output[j]

            new_batches = (min, max)

            return new_batches

    return RebatchWrapper()

if __name__ == '__main__':
    # add a test here

    hklin = os.path.join(os.environ['XIA2_ROOT'],
                         'Data', 'Test', 'Mtz', '12287_1_E1_1_10.mtz')
    hklout = 'temp.mtz'

    rb = Rebatch()

    rb.set_hklin(hklin)
    rb.set_hklout(hklout)
    rb.set_first_batch(100)

    new = rb.rebatch()

    print new