from __future__ import annotations

import sys

import numpy as np

from cctbx import factor_kev_angstrom
from cctbx.eltbx import attenuation_coefficient
from libtbx.phil import parse
from scitbx.array_family import flex
from scitbx.matrix import col

from dxtbx.format.FormatXTC import FormatXTC, locator_str
from dxtbx.model import Detector, ParallaxCorrectedPxMmStrategy

epix_locator_str = """
"""

epix_locator_scope = parse(epix_locator_str + locator_str, process_includes=True)


class FormatXTCEpix(FormatXTC):
    def __init__(self, image_file, **kwargs):
        super().__init__(image_file, locator_scope=epix_locator_scope, **kwargs)
        self._cached_detector = {}

    @staticmethod
    def understand(image_file):
        try:
            params = FormatXTC.params_from_phil(epix_locator_scope, image_file)
        except Exception:
            return False
        return any("epix" in src.lower() for src in params.detector_address)

    def get_raw_data(self, index=None):
        if index is None:
            index = 0
        d = FormatXTCEpix.get_detector(self, index)
        evt = self._get_event(index)
        run = self.get_run_from_index(index)
        det = self._get_psana_detector(run)
        try:
            data = det.calib(evt)
        except Exception:
            # psana2
            data = det.raw.calib(evt)
        data = data.astype(np.float64)
        # the shape of the epix 10k data is (16, 352, 384)

        self._raw_data = []
        for module_count, module in enumerate(d.hierarchy()):
            for asic_count, asic in enumerate(module):
                fdim, sdim = asic.get_image_size()
                sensor_id = asic_count // 2  # There are 2X2 asics per module
                asic_in_sensor_id = asic_count % 2  # this number will be 0 or 1
                asic_data = data[module_count][
                    sensor_id * sdim : (sensor_id + 1) * sdim,
                    asic_in_sensor_id * fdim : (asic_in_sensor_id + 1) * fdim,
                ]  # 8 sensors per module
                self._raw_data.append(flex.double(np.array(asic_data)))
        assert len(d) == len(self._raw_data), (len(d), len(self._raw_data))
        return tuple(self._raw_data)

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

    def _detector(self, index=None):
        try:
            from PSCalib.SegGeometryStore import sgs
        except ModuleNotFoundError:
            # psana2
            from psana.pscalib.geometry.SegGeometryStore import sgs

        from serialtbx.detector.xtc import basis_from_geo

        run = self.get_run_from_index(index)
        try:
            try:
                run_num = run.run()
            except Exception:
                # psana2_idx mode
                run_num = run.runnum
        except AttributeError:
            # smd and psana2 modes
            run_num = run
        if run_num in self._cached_detector:
            return self._cached_detector[run_num]

        if index is None:
            index = 0
        assert len(self.params.detector_address) == 1

        wavelength = self.get_beam(index).get_wavelength()

        det = self._get_psana_detector(run)

        try:
            geom = det.pyda.geoaccess(self._get_event(index).run())
        except AttributeError:
            # psana2
            geom = det.raw._det_geo()
        try:
            pixel_size_um = det.pixel_size(self._get_event(index))
        except AttributeError:
            pixel_size_um = geom.get_pixel_scale_size()
        pixel_size = pixel_size_um / 1000.0  # convert to mm
        d = Detector()
        pg0 = d.hierarchy()
        # first deal with D0
        det_num = 0
        root = geom.get_top_geo()
        root_basis = basis_from_geo(root)
        while len(root.get_list_of_children()) == 1:
            sub = root.get_list_of_children()[0]
            sub_basis = basis_from_geo(sub)
            root = sub
            root_basis = root_basis * sub_basis
        t = root_basis.translation
        root_basis.translation = col((t[0], t[1], -abs(t[2])))

        origin = col((root_basis * col((0, 0, 0, 1)))[0:3])
        fast = col((root_basis * col((1, 0, 0, 1)))[0:3]) - origin
        slow = col((root_basis * col((0, 1, 0, 1)))[0:3]) - origin
        pg0.set_local_frame(fast.elems, slow.elems, origin.elems)
        pg0.set_name("D%d" % (det_num))

        # Modules next
        for module_num in range(len(root.get_list_of_children())):
            pg1 = pg0.add_group()
            module = root.get_list_of_children()[module_num]
            module_basis = basis_from_geo(module)
            origin = col((module_basis * col((0, 0, 0, 1)))[0:3])
            fast = col((module_basis * col((1, 0, 0, 1)))[0:3]) - origin
            slow = col((module_basis * col((0, 1, 0, 1)))[0:3]) - origin
            pg1.set_local_frame(fast.elems, slow.elems, origin.elems)
            pg1.set_name("D%dM%d" % (det_num, module_num))

            # Read the known layout of the Epix 2x2 module
            sg = sgs.Create(segname=module.oname)
            xx, yy = sg.get_seg_xy_maps_um()
            xx = xx / 1000
            yy = yy / 1000

            # Now deal with ASICs
            for asic_num in range(4):
                val = "ARRAY_D0M%dA%d" % (module_num, asic_num)
                p = pg1.add_panel()
                dim_slow = xx.shape[0]
                dim_fast = xx.shape[1]
                sensor_id = asic_num // 2  # There are 2X2 asics per module
                asic_in_sensor_id = asic_num % 2  # this number will be 0 or 1
                id_slow = sensor_id * (dim_slow // 2)
                id_fast = asic_in_sensor_id * (dim_fast // 2)
                origin = col((xx[id_slow][id_fast], yy[id_slow][id_fast], 0))
                fp = col((xx[id_slow][id_fast + 1], yy[id_slow][id_fast + 1], 0))
                sp = col((xx[id_slow + 1][id_fast], yy[id_slow + 1][id_fast], 0))
                fast = (fp - origin).normalize()
                slow = (sp - origin).normalize()
                p.set_local_frame(fast.elems, slow.elems, origin.elems)
                p.set_pixel_size((pixel_size, pixel_size))
                p.set_image_size((dim_fast // 2, dim_slow // 2))
                p.set_trusted_range((-100, 2e6))
                p.set_gain(factor_kev_angstrom / wavelength)
                p.set_name(val)

        try:
            beam = self._beam(index)
        except Exception:
            print(
                "No beam object initialized. Returning ePix detector without parallax corrections"
            )
            try:
                run_num = run.run()
            except AttributeError:
                # smd and psana2 modes
                run_num = run
            if run_num in self._cached_detector:
                self._cached_detector[run_num] = d
                # self._cached_detector[run.run()] = d
            return d

        # take into consideration here the thickness of the sensor also the
        # wavelength of the radiation (which we have in the same file...)
        wavelength = beam.get_wavelength()
        thickness = 0.5  # mm

        material = "Si"
        table = attenuation_coefficient.get_table(material)
        # mu_at_angstrom returns cm^-1
        mu = table.mu_at_angstrom(wavelength) / 10.0  # mu: mm^-1
        t0 = thickness
        for panel in d:
            panel.set_px_mm_strategy(ParallaxCorrectedPxMmStrategy(mu, t0))
            panel.set_thickness(thickness)
            panel.set_material(material)
            panel.set_mu(mu)

        try:
            run_num = run.run()
        except AttributeError:
            # smd and psana2 modes
            run_num = run
        self._cached_detector[run_num] = d
        return d


if __name__ == "__main__":
    for arg in sys.argv[1:]:
        # Bug, should call this part differently for understand method to work
        print(FormatXTCEpix.understand(arg))
