"""
Experimental format for the X-Spectrum LAMBDA detector
http://www.x-spectrum.de/
"""

from __future__ import annotations

import sys

import h5py

from cctbx.eltbx import attenuation_coefficient
from scitbx import matrix
from scitbx.array_family import flex

from dxtbx.format.Format import Format
from dxtbx.format.FormatHDF5 import FormatHDF5
from dxtbx.model import ParallaxCorrectedPxMmStrategy


class FormatHDF5Lambda(FormatHDF5):
    @staticmethod
    def understand(image_file):
        try:
            tag = FormatHDF5.open_file(image_file, "rb").read(8)
        except OSError:
            return False

        # check that this is a HDF5 file (should not have got here if not
        # anyway...)

        if tag != b"\211HDF\r\n\032\n":
            return False

        with h5py.File(image_file, "r") as h5_handle:
            try:
                desc = h5_handle["entry/instrument/detector/description"]
            except KeyError:
                return False

            try:
                if b"Lambda" in desc[()][0]:
                    return True
            except TypeError:
                return False

        return False

    def _start(self):
        self._h5_handle = h5py.File(self.get_image_file(), "r")

    def _goniometer(self):
        """Dummy goniometer - assume vertical (EMBl-HH P14)"""

        return self._goniometer_factory.known_axis((0, 1, 0))

    def _detector(self):
        """Partly dummy detector"""
        # Get the detector geometry
        entry = self._h5_handle["entry"]
        instrument = entry["instrument"]
        detector = instrument["detector"]

        # Initialise detector frame - origin at 0,0,0
        fast = matrix.col((1.0, 0.0, 0.0))
        slow = matrix.col((0.0, 1.0, 0.0))
        orig = matrix.col((0.0, 0.0, 0.0))

        # Get the pixel and image size
        pixel_size = (
            1.0e-3 * float(detector["x_pixel_size"][0]),
            1.0e-3 * float(detector["y_pixel_size"][0]),
        )
        layout = detector["layout"][0].split(b"X")
        image_size = int(layout[0]), int(layout[1])
        trusted_range = (0, int(detector["saturation_value"][0]))
        thickness = float(detector["sensor_thickness"][0]) / 1000.0
        material = detector["sensor_material"][0].decode()

        # Make the detector
        detector = self._detector_factory.make_detector(
            "PAD",
            fast,
            slow,
            orig,
            pixel_size,
            image_size,
            trusted_range,
            name="Panel",
            thickness=thickness,
            material=material,
        )

        # At the moment, beam is a dummy object because wavelength is not set in
        # the header. Therefore, the px<-->mm strategy will generally be
        # incorrect. Set it anyway, to override later.
        beam = self._beam()
        wavelength = beam.get_wavelength()

        # this will fail for undefined composite materials
        table = attenuation_coefficient.get_table(material)

        # mu_at_angstrom returns cm^-1, but need mu in mm^-1
        mu = table.mu_at_angstrom(wavelength) / 10.0

        for panel in detector:
            panel.set_mu(mu)
            panel.set_px_mm_strategy(ParallaxCorrectedPxMmStrategy(mu, thickness))

        return detector

    def _beam(self):
        """Dummy beam"""

        wavelength = 1.0
        return self._beam_factory.simple(wavelength)

    def _scan(self):
        """Dummy scan"""

        entry = self._h5_handle["entry"]
        nframes = int(entry["instrument/detector/collection/frame_numbers"][0])
        image_range = (1, nframes)
        exposure_times = 0.0
        oscillation = (0, 1)
        epochs = [0] * nframes

        return self._scan_factory.make_scan(
            image_range, exposure_times, oscillation, epochs, deg=True
        )

    def get_num_images(self):
        detector = self._h5_handle["entry/instrument/detector"]
        data = detector["data"]
        return data.shape[0]

    def get_goniometer(self, index=None):
        return Format.get_goniometer(self)

    def get_detector(self, index=None):
        return Format.get_detector(self)

    def get_beam(self, index=None):
        return Format.get_beam(self)

    def get_scan(self, index=None):
        if index is None:
            return Format.get_scan(self)
        else:
            scan = Format.get_scan(self)
            return scan[index]

    def get_raw_data(self, index):
        detector = self._h5_handle["entry/instrument/detector"]
        data = detector["data"]
        im = data[index, :, :].astype("int32")  # convert from int16
        return flex.int(im)

    def get_image_file(self, index=None):
        return Format.get_image_file(self)


if __name__ == "__main__":
    for arg in sys.argv[1:]:
        print(FormatHDF5Lambda.understand(arg))
