#!/usr/bin/env python
# LabelitDistl.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.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

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 LabelitDistl(DriverType = None):
    '''Factory for LabelitDistl wrapper classes, with the specified
    Driver type.'''

    DriverInstance = DriverFactory.Driver(DriverType)

    class LabelitDistlWrapper(DriverInstance.__class__):
        '''A wrapper for the program labelit.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.distl')

            self._images = []

            self._statistics = { }

        def add_image(self, image):
            '''Add an image for indexing.'''

            if not image in self._images:
                self._images.append(image)

            return

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

            self._images.sort()

            for i in self._images:
                self.add_command_line(i)

            task = 'Screen images:'

            for i in self._images:
                task += ' %s' % i

            self.set_task(task)

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

if __name__ == '__main__':

    # run a demo test

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

    l = LabelitDistl()
    for image in sys.argv[1:]:
        l.add_image(image)

    l.distl()

    for image in sys.argv[1:]:
        stats = l.get_statistics(image)

        print image, stats['spots_good'], stats['spots']