#!/usr/bin/env python
# DistlSweepStrength.py
#   Copyright (C) 2013 Diamond Light Source, Richard Gildea
#
#   This code is distributed under the BSD license, a copy of which is
#   included in the root directory of this package.
#
# 11th June 2013
#
# A wrapper for the replacement for distl.sweep_strength.

import libtbx.load_env

import os
import sys

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

sys.path.append(os.path.join(os.environ['XIA2CORE_ROOT'],
                             'Python'))

from Driver.DriverFactory import DriverFactory

import spotfinder
from spotfinder.command_line import sweep_strength
import libtbx.phil

master_params = sweep_strength.master_params
custom_params = libtbx.phil.parse("""
distl.verbosity = 0
""")
# override distl.sweep_strength_defaults
master_params = master_params.fetch(source=custom_params)


def DistlSweepStrength(DriverType=None, params=None):
  '''Factory for DistlSweepStrength wrapper classes, with the specified
  Driver type.'''

  DriverInstance = DriverFactory.Driver(DriverType)

  class DistlSweepStrengthWrapper(DriverInstance.__class__):
    '''A wrapper for the program distl.sweep_strength - which will provide
    functionality for looking for finding spots &c.'''

    def __init__(self, params=None):

      DriverInstance.__class__.__init__(self)

      # phil parameters

      if not params:
        params = master_params.extract()
      self._params = params



      self.set_executable('distl.sweep_strength')

      self._images = []

      self._statistics = { }
      self._peaks = []

      return

    def set_image(self, image):
      '''Set an image for analysis.'''

      self._images.append(image)

      return

    def distl(self):
      '''Actually analyse the images.'''

      self.add_command_line(self._image)
      self.start()
      self.close_wait()

      # check for errors
      self.check_for_errors()

      # ok now we're done, let's look through for some useful stuff

      output = self.get_all_output()

      for o in output:
        l = o.split()

        if l[:2] == ['Spot', 'Total']:
          self._statistics['spots_total'] = int(l[-1])
        if l[:2] == ['In-Resolution', 'Total']:
          self._statistics['spots'] = int(l[-1])
        if l[:3] == ['Good', 'Bragg', 'Candidates']:
          self._statistics['spots_good'] = int(l[-1])
        if l[:2] == ['Ice', 'Rings']:
          self._statistics['ice_rings'] = int(l[-1])
        if l[:3] == ['Method', '1', 'Resolution']:
          self._statistics['resol_one'] = float(l[-1])
        if l[:3] == ['Method', '2', 'Resolution']:
          self._statistics['resol_two'] = float(l[-1])
        if l[:3] == ['%Saturation,', 'Top', '50']:
          self._statistics[
              'saturation'] = float(l[-1])

      return

    def run(self):
      '''Actually analyse the images.'''

      for image in self._images:
        self.add_command_line(image)
      self.start()
      self.close_wait()

      # check for errors
      self.check_for_errors()

      output = self.get_all_output()
      print "".join(output)

  return DistlSweepStrengthWrapper()

if __name__ == '__main__':

  if len(sys.argv) < 2:
    raise RuntimeError, '%s image' % sys.argv[0]

  print '%s ... %s' % (sys.argv[1], sys.argv[-1])

  d = DistlSweepStrength()
  for arg in sys.argv[1:]:
    if os.path.isfile(arg):
      d.set_image(arg)
    else:
      d.add_command_line(arg)
  d.run()