#!/usr/bin/env python
# Sortmtz.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.
#
# 5th June 2006
#
# A wrapper for the CCP4 program sortmtz.
#
#
# Provides:
#
# Reflection sorting functionality for going from Mosflm to Scala.
#

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'))

if not os.environ['XIA2_ROOT'] in sys.path:
  sys.path.append(os.environ['XIA2_ROOT'])

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

def Sortmtz(DriverType = None):
  '''A factory for SortmtzWrapper classes.'''

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

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

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

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

      self._sort_order = 'H K L M/ISYM BATCH'

      self._hklin_files = []

      return

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

    def check_sortmtz_errors(self):
      '''Check the output for "standard" errors.'''

      lwbat_warning = ''

      for l in self.get_all_output():

        if 'From ccp4_lwbat: warning:' in l:
          lwbat_warning = l.split('warning:')[1].strip()

        if 'error in ccp4_lwbat' in l:
          raise RuntimeError, lwbat_warning

        if 'Sorting failed' in l:
          raise RuntimeError, 'sorting failed'

        if 'Inconsistent operator orders in input file' in l:
          raise RuntimeError, 'different sort orders'

    def sort(self, vrset = None):
      '''Actually sort the reflections.'''

      # if we have not specified > 1 hklin file via the add method,
      # check that the set_hklin method has been used.
      if not self._hklin_files:
        self.check_hklin()

      self.check_hklout()

      if self._hklin_files:
        task = ''
        for hklin in self._hklin_files:
          task += ' %s' % hklin
        self.set_task('Sorting reflections%s => %s' % \
                     (task,
                      os.path.split(self.get_hklout())[-1]))
      else:
        self.set_task('Sorting reflections %s => %s' % \
                     (os.path.split(self.get_hklin())[-1],
                      os.path.split(self.get_hklout())[-1]))

      self.start()

      # allow for the fact that large negative reflections may
      # result from XDS output...

      if vrset:
        self.input('VRSET_MAGIC %f' % vrset)

      self.input(self._sort_order)

      # multiple mtz files get passed in on the command line...

      if self._hklin_files:
        for m in self._hklin_files:
          # FIXME have removed the quotes as this breaks the parser
          # it is assumed that the file is a comment! This is not
          # a problem in 6.0.2 - so allow for the quotes to
          # be added if and only if there is a space...
          if ' ' in m:
            Chatter.write(
                'Quoting input files - you have been warned!')
            self.input('"%s"' % m)
          else:
            self.input('%s' % m)

      self.close_wait()

      try:

        # general errors - SEGV and the like
        self.check_for_errors()

        # ccp4 specific errors
        self.check_ccp4_errors()
        if 'Error' in self.get_ccp4_status():
          raise RuntimeError, '[SORTMTZ] %s' % status

        # sortmtz specific errors
        self.check_sortmtz_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 SortmtzWrapper()

if __name__ == '__main_2_':
  # run some tests

  if not os.environ.has_key('XIA2CORE_ROOT'):
    raise RuntimeError, 'XIA2CORE_ROOT not defined'

  dpa = os.environ['XIA2_ROOT']

  hklin1 = os.path.join(dpa,
                        'Data', 'Test', 'Mtz', '12287_1_E1_1_10.mtz')
  hklin2 = os.path.join(dpa,
                        'Data', 'Test', 'Mtz', '12287_1_E1_11_20.mtz')

  s = Sortmtz()
  s.add_hklin(hklin1)
  s.add_hklin(hklin2)
  s.set_hklout('null.mtz')

  try:
    print s.sort()
  except RuntimeError, e:
    print 'Error => %s' % e

  s = Sortmtz()
  s.add_hklin(hklin1)
  s.add_hklin(hklin1)
  s.set_hklout('null.mtz')

  try:
    print s.sort()
  except RuntimeError, e:
    print 'Error => %s' % e

if __name__ == '__main__':
  s = Sortmtz()
  s.add_hklin('TS00_13185_unmerged_INFL.mtz')
  s.set_hklout('TS00_13185_sorted_INFL.mtz')
  s.sort()