#!/usr/bin/env python
# LabelitDistl.py
#   Copyright (C) 2010 Diamond Light Source, Graeme Winter
#
#   This code is distributed under the BSD license, a copy of which is
#   included in the root directory of this package.
#
# 11th May 2010
#
# A wrapper for the replacement for labelit.distl - distl.signal_strength.
# This includes the added ability to get a list of the spot positions on
# the image. This can in turn replace printpeaks.
#
# N.B. this is only included in more recent versions of Labelit.

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

    DriverInstance = DriverFactory.Driver(DriverType)

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

        def __init__(self):

            DriverInstance.__class__.__init__(self)

            self.set_executable('distl.signal_strength')

            self._image = None

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

            return

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

            self._image = 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 find_peaks(self):
            '''Actually analyse the images.'''

            self.add_command_line(self._image)
            self.add_command_line('verbose=True')
            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()

            self._peaks = []

            peak_xy = None
            peak_sn = None

            for o in output:

                if not 'Peak' in o:
                    continue

                if 'signal-to-noise' in o:
                    assert(peak_xy is None)
                    peak_sn = float(o.split('=')[-1])

                elif 'position' in o:
                    assert(peak_sn)
                    l = o.replace('=', ' ').split()
                    peak_xy = float(l[3]), float(l[5])

                    # it appears that the printpeaks coordinate system is
                    # swapped w.r.t. the distl one...

                    self._peaks.append((peak_xy[1], peak_xy[0], peak_sn))

                    peak_xy = None
                    peak_sn = None

            return self._peaks

        def get_statistics(self):
            return self._statistics

        def get_peaks(self):
            return self._peaks

    return DistlSignalStrengthWrapper()

if __name__ == '__main__':

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

    d = DistlSignalStrength()
    d.set_image(sys.argv[1])
    peaks = d.find_peaks()

    for m in peaks:
        print '%6.1f %6.1f %6.1f' % m