#!/usr/bin/env python
# LabelitStats_distl.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.
#
# 2nd June 2006
#
# A wrapper for labelit.stats_distl - this will provide functionality to:
#
# Looking for ice rings.
# Screening the images.
#
# The output looks like:
#
#                     File : 12287_1_E1_001.img
#               Spot Total :   1568
#      In-Resolution Total :   1421
#    Good Bragg Candidates :   1135
#                Ice Rings :      2
#      Method 1 Resolution :   1.78
#      Method 2 Resolution :   1.90
#        Maximum unit cell :  242.8
#%Saturation, Top 50 Peaks :  14.35
#
# Based on the output of the last Distl run.
#
# The output should be identical to the last labelit.distl run.
# Note that this does not take any image as input

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

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

  DriverInstance = DriverFactory.Driver(DriverType)

  class LabelitStats_distlWrapper(DriverInstance.__class__):
    '''A wrapper for the program labelit.stats_distl - which will provide
    functionality for looking for ice rings and screening diffraction
    images.'''

    def __init__(self):

      DriverInstance.__class__.__init__(self)

      self.set_executable('labelit.stats_distl')

      self._statistics = { }

    def stats_distl(self):
      '''Return the image statistics.'''

      self.set_task('Return statistics from last distl run')

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

      current_image = None

      for o in output:

        if 'None' in o and 'Resolution' in o:
          l = o.replace('None', '0.0').split()
        else:
          l = o.split()

        if l[:1] == ['File']:
          current_image = l[2]
          self._statistics[current_image] = { }

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

      return 'ok'

    # things to get results from the indexing

    def get_statistics(self, image):
      '''Get the screening statistics from image as dictionary.
      The keys are spots_total, spots, spots_good, ice_rings,
      resol_one, resol_two.'''

      return self._statistics[os.path.split(image)[-1]]

  return LabelitStats_distlWrapper()

if __name__ == '__main__':

  # run a demo test

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

  l = LabelitStats_distl()
  l.stats_distl()

  stats = l.get_statistics('12287_1_E1_001.img')

  print 'Fraction of good spots: %4.2f' % (float(stats['spots_good']) /
                                           float(stats['spots']))
  print 'Ice rings:              %d' % stats['ice_rings']
  print 'Resolutions:            %f  %f' % \
        (stats['resol_one'], stats['resol_two'])