#!/usr/bin/env python
# FormatTIFFRayonix.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.
#
# An implementation of the TIFF image reader for Rayonix images. Inherits from
# FormatTIFF.

import time
import datetime
import struct

from dxtbx.format.FormatTIFF import FormatTIFF

class FormatTIFFRayonix(FormatTIFF):
    '''A class for reading TIFF format Rayonix images, and correctly
    constructing a model for the experiment from this.'''

    @staticmethod
    def understand(image_file):
        '''Check to see if this looks like an Rayonix TIFF format image,
        i.e. we can make sense of it. This simply checks that records which
        describe the size of the image match with the TIFF records which do
        the same.'''

        if FormatTIFF.understand(image_file) == 0:
            return 0

        width, height, depth, order, bytes = FormatTIFF.get_tiff_header(
            image_file)

        assert(len(bytes) == 4096)

        if order == FormatTIFF.LITTLE_ENDIAN:
            endian = '<'
        else:
            endian = '>'

        _I = endian + 'I'
        _i = endian + 'i'

        _width = struct.unpack(_I, bytes[1024 + 80:1024 + 84])[0]
        _height = struct.unpack(_I, bytes[1024 + 84:1024 + 88])[0]
        _depth = struct.unpack(_I, bytes[1024 + 88:1024 + 92])[0]

        if width != _width or height != _height or depth != _depth:
            return 0

        nimages = struct.unpack(_I, bytes[1024 + 112:1024 + 116])[0]
        origin = struct.unpack(_I, bytes[1024 + 116:1024 + 120])[0]
        orientation = struct.unpack(_I, bytes[1024 + 120:1024 + 124])[0]
        view = struct.unpack(_I, bytes[1024 + 124:1024 + 128])[0]

        if nimages != 1 or origin != 0 or orientation != 0 or view != 0:
            return 0

        return 2

    def __init__(self, image_file):
        '''Initialise the image structure from the given file, including a
        proper model of the experiment.'''

        assert(FormatTIFFRayonix.understand(image_file) > 0)

        width, height, depth, order, bytes = FormatTIFF.get_tiff_header(
            image_file)

        if order == FormatTIFF.LITTLE_ENDIAN:
            self._I = '<I'
            self._i = '<i'
            self._ii = '<ii'
        else:
            self._I = '>I'
            self._i = '>i'
            self._ii = '>ii'

        FormatTIFF.__init__(self, image_file)

        return

    # FIXME have implemented none of those which follow...

    def _goniometer(self):
        '''Return a model for goniometer corresponding to the values stored
        in the image header. In the first instance assume this is a single
        axis annd raise exception otherwise.'''

        starts, ends, offset, width = self._get_rayonix_scan_angles()

        for j in range(len(starts)):
            if j != offset:
                assert(starts[j] == 0.0)
                assert(ends[j] == 0.0)

        return self._goniometer_factory.single_axis()

    def _detector(self):
        '''Return a model for a simple detector, which at the moment insists
        that the offsets and rotations are all 0.0.'''

        starts, ends, offset, width = self._get_rayonix_scan_angles()
        rotations = self._get_rayonix_detector_rotations()

        # assert that two-theta offset is 0.0

        assert(starts[0] == 0.0)
        assert(ends[0] == 0.0)

        # assert that the rotations are all 0.0

        assert(rotations[0] == 0.0)
        assert(rotations[1] == 0.0)
        assert(rotations[2] == 0.0)

        distance = self._get_rayonix_distance()
        beam_x, beam_y = self._get_rayonix_beam_xy()
        pixel_size = self._get_rayonix_pixel_size()
        image_size = self._tiff_width, self._tiff_height
        overload = struct.unpack(
            self._i, self._tiff_header_bytes[1128:1132])[0]
        underload = 0

        beam = beam_x * pixel_size[0], beam_y * pixel_size[1]

        return self._detector_factory.simple(
            'CCD', distance, beam, '+x', '-y', pixel_size,
            image_size, (underload, overload), [])

    def _beam(self):
        '''Return a simple model for the beam.'''

        wavelength = struct.unpack(
            self._i, self._tiff_header_bytes[1932:1936])[0] * 1.0e-5

        return self._beam_factory.simple(wavelength)

    def _scan(self):
        '''Return the scan information for this image.'''

        format = self._scan_factory.format('TIFF')
        exposure_time = self._get_rayonix_times()[1]
        epoch = time.mktime(self._get_rayonix_timestamp())

        starts, ends, offset, width = self._get_rayonix_scan_angles()

        osc_start = starts[offset]
        osc_range = width

        return self._scan_factory.single(
            self._image_file, format, exposure_time,
            osc_start, osc_range, epoch)

    ####################################################################
    #                                                                  #
    # Helper methods to get all of the values out of the TIFF header   #
    # - separated out to assist with code clarity                      #
    #                                                                  #
    ####################################################################

    def _get_rayonix_distance(self):
        '''Look in the usual places for the detector distance, return this
        as a float in mm.'''

        distance = struct.unpack(
            self._i, self._tiff_header_bytes[1664:1668])[0]

        if distance != 0:
            return distance * 0.001

        distance = struct.unpack(
            self._i, self._tiff_header_bytes[1720:1724])[0]

        if distance != 0:
            return distance * 0.001

        raise RuntimeError, 'cannot find distance in header'

    def _get_rayonix_beam_xy(self):
        '''Get the beam x, y positions which are defined in the standard
        to be in pixels. X and Y are not defined by the documentation, so
        taking as a given that these are horizontal and vertical. N.B.
        the documentation states that the horizontal direction is fast.'''

        beam_x, beam_y = struct.unpack(
            self._ii, self._tiff_header_bytes[1668:1676])[:2]

        return beam_x * 0.001, beam_y * 0.001

    def _get_rayonix_pixel_size(self):
        '''Get the pixel sizes in mm.'''

        pixel_x, pixel_y = struct.unpack(
            self._ii, self._tiff_header_bytes[1796:1804])[:2]

        return pixel_x * 1.0e-6, pixel_y * 1.0e-6

    def _get_rayonix_times(self):
        '''Get the integration, exposure times in seconds.'''

        integration, exposure = struct.unpack(
            self._ii, self._tiff_header_bytes[1676:1684])[:2]

        return integration * 0.001, exposure * 0.001

    def _get_rayonix_timestamp(self):
        '''Get the image acquisition timestamp.'''

        timestamp = self._tiff_header_bytes[2368:2400]

        month = int(timestamp[:2])
        day = int(timestamp[2:4])
        hour = int(timestamp[4:6])
        minute = int(timestamp[6:8])
        year = int(timestamp[8:12])
        second = int(timestamp[13:15])
        return datetime.datetime(year, month, day,
                                 hour, minute, second).timetuple()

    def _get_rayonix_scan_angles(self):
        '''Get the scan angles for: twotheta, omega, chi, kappa, phi, delta,
        gamma. The exact definitions for these are somewhat poorly defined,
        though I presume that they will come from some kind of standard
        goniometer definition... Also returns the scan axis offset and
        the apparent oscillation width, as these are sometimes not properly
        recorded elsewhere.'''

        if self._tiff_byte_order == FormatTIFF.LITTLE_ENDIAN:
            iiiiiii = '<iiiiiii'
        else:
            iiiiiii = '>iiiiiii'

        start_angles = struct.unpack(
            iiiiiii, self._tiff_header_bytes[1692:1720])[:7]

        end_angles = struct.unpack(
            iiiiiii, self._tiff_header_bytes[1724:1752])[:7]

        axis_offset = struct.unpack(
            self._i, self._tiff_header_bytes[1756:1760])[0]

        axis_range = struct.unpack(
            self._i, self._tiff_header_bytes[1760:1764])[0]

        starts_degrees = [s * 0.001 for s in start_angles]
        ends_degrees = [e * 0.001 for e in end_angles]

        return starts_degrees, ends_degrees, axis_offset, axis_range * 0.001

    def _get_rayonix_detector_rotations(self):
        '''Get the recorded rotx, roty, rotz of the detector - which in most
        cases will probably all be 0.0.'''

        if self._tiff_byte_order == FormatTIFF.LITTLE_ENDIAN:
            iii = '<iii'
        else:
            iii = '>iii'

        rot_angles = struct.unpack(
            iii, self._tiff_header_bytes[1764:1776])[:3]

        rot_degrees = [r * 0.001 for r in rot_angles]

        return rot_degrees

if __name__ == '__main__':

    import sys

    for arg in sys.argv[1:]:
        print FormatTIFFRayonix.understand(arg)