#!/usr/bin/env python
# PrintpeaksMosflm.py
#   Copyright (C) 2011 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.
#
# A replacement for the printpeaks tool and wrapper, using Mosflm.
#

import os
import sys
import copy
import math

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

if not os.environ.has_key('XIA2_ROOT'):
    raise RuntimeError, 'XIA2_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 Diffdump import Diffdump

from Experts.FindImages import image2template_directory, image2image, \
     template_number2image

def PrintpeaksMosflm(DriverType = None):
    '''A factory for wrappers for the printpeaks/mosflm.'''

    DriverInstance = DriverFactory.Driver(DriverType)

    class PrintpeaksMosflmWrapper(DriverInstance.__class__):
        '''Provide access to the spot finding functionality in Mosflm.'''

        def __init__(self):
            DriverInstance.__class__.__init__(self)

            self.set_executable('ipmosflm')

            if 'CCP4_SCR' in os.environ:
                self.set_working_directory(os.environ['CCP4_SCR'])

            self._image = None
            self._peaks = { }

            return

        def set_image(self, image):
            '''Set an image to read the header of.'''
            self._image = image
            self._peaks = { }
            return

        def get_maxima(self, threshold = 5.0):
            '''Run diffdump, printpeaks to get a list of diffraction maxima
            at their image positions, to allow for further analysis.'''

            if not self._image:
                raise RuntimeError, 'image not set'

            if not os.path.exists(self._image):
                raise RuntimeError, 'image %s does not exist' % \
                      self._image

            dd = Diffdump()
            dd.set_image(self._image)
            header = dd.readheader()

            beam = header['raw_beam']
            pixel = header['pixel']

            template, directory = image2template_directory(self._image)
            image_number = image2image(os.path.split(self._image)[-1])

            spot_file = '%s.spt' % template_number2image(
                template, image_number)

            self.start()

            self.input('template "%s"' % template)
            self.input('directory "%s"' % directory)
            self.input('findspots local find %d file %s' % \
                       (image_number, spot_file))
            self.input('go')

            self.close_wait()

            self.check_for_errors()

            output = open(os.path.join(self.get_working_directory(),
                                       spot_file)).readlines()

            os.remove(os.path.join(self.get_working_directory(),
                                   spot_file))

            peaks = []

            for record in output[3:-2]:
                lst = record.split()
                x = float(lst[0])
                y = float(lst[1])
                i = float(lst[4]) / float(lst[5])
                x /= pixel[0]
                y /= pixel[1]

                if i < threshold:
                    continue

                # this is Mosflm right? Swap X & Y!!

                peaks.append((y, x, i))

            return peaks

        def printpeaks(self):
            '''Run printpeaks and get the list of peaks out, then decompose
            this to a histogram.'''

            if not self._image:
                raise RuntimeError, 'image not set'

            if not os.path.exists(self._image):
                raise RuntimeError, 'image %s does not exist' % \
                      self._image

            _peaks = self.get_maxima()
            peaks = []

            for peak in _peaks:
                intensity = peak[2]
                peaks.append(intensity)

            # now construct the histogram

            log_max = int(math.log10(peaks[0])) + 1
            max_limit = int(math.pow(10.0, log_max))

            if False:

                limit = math.pow(10.0, log_max)

                while limit > 2.0:
                    self._peaks[limit] = len([j for j in peaks if j > limit])
                    limit *= 0.5

            else:

                for limit in [5, 10, 20, 50, 100, 200, 500, 1000]:
                    if limit > max_limit:
                        continue
                    self._peaks[float(limit)] = len(
                        [j for j in peaks if j > limit])


            return self._peaks

        def threshold(self, nspots):
            if not self._peaks:
                peaks = self.printpeaks()
            else:
                peaks = self._peaks
            keys = peaks.keys()
            keys.sort()
            keys.reverse()
            for thresh in keys:
                if peaks[thresh] > nspots:
                    return thresh
            return min(keys)

        def screen(self):
            if not self._image:
                raise RuntimeError, 'image not set'

            if not os.path.exists(self._image):
                raise RuntimeError, 'image %s does not exist' % \
                      self._image

            _peaks = self.get_maxima()

            peaks = []

            for peak in _peaks:
                if peak[2] > 10:
                    peaks.append(peak)

            if len(peaks) < 10:
                return 'blank'

            return 'ok'

    return PrintpeaksMosflmWrapper()

if __name__ == '__main-old__':

    import time

    def printer(peaks):
        keys = peaks.keys()
        keys.sort()
        for k in keys:
            print '%.5f %d' % (k, peaks[k])

    directory = os.path.join(os.environ['XIA2_ROOT'],
                             'Data', 'Test', 'Images')

    if len(sys.argv) == 1:
        p = Printpeaks()
        image = os.path.join(directory, '12287_1_E1_001.img')
        p.set_image(image)
        peaks = p.printpeaks()
        printer(peaks)

    else:

        # for image in sys.argv[1:]:
        if 1 == 0:

            print image
            p = Printpeaks()
            p.set_image(image)

            peaks = p.printpeaks()
            printer(peaks)

            thresh = p.threshold(200)
            print '200 peak threshold: %f' % thresh

        t0 = time.time()
        count = 0
        for image in sys.argv[1:]:
            count += 1
            p = Printpeaks()
            p.set_image(image)
            status = p.screen()
            print os.path.split(image)[-1], status
        t1 = time.time()

        print 'Total time: %.1f' % (t1 - t0)
        print 'Per image: %.3f' % ((t1 - t0) / count)

if __name__ == '__main__':
    # run a test of some of the new code...

    p = Printpeaks()
    p.set_image(sys.argv[1])
    peaks = p.get_maxima()

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