from __future__ import annotations

import os
import pint
import numpy
import h5py

from nxtomo.application.nxtomo import NXtomo
from nxtomo.nxobject.nxdetector import ImageKey

from nxtomomill.tests.datasets import GitlabDataset
from nxtomomill import converter
from nxtomomill.models.h52nx import H52nxModel
from nxtomomill.models.h52nx.keys_section import _cast_keys_section_to_tuple

_ureg = pint.get_application_registry()


def test_holotomo_conversion_id16a(tmp_path):
    """
    test conversion on one id16a use case

    check x, y and z positions
    warning: bliss-tomo is using esrf coordinate system when
    NXtomo is using the McStas one. see https://tomotools.gitlab-pages.esrf.fr/tomwer/development/referential.html for more details.
    on esrf referential we expect the following (from peter):

        translation_x = sx (to keep things simple)
        translation_y = sy + spy
        translation_z = sz + spz
        sample_x = sax
        sample_y = say

    .. warning:: In 'tomoconfig'

    'translation_y is defined as 'sy' instead of sy + spy
    'translation_z is defined as 'sz' instead of sz + spz

    So this is why we need to overwrite the settings
    """
    dataset_file = os.path.join(
        "h5_datasets", "holotomo", "id16a", "SiemensStar_HT_025nm_0002.h5"
    )
    input_file = GitlabDataset.get_dataset(dataset_file)
    assert os.path.exists(input_file)
    output_file = os.path.join(tmp_path, "nexus_file.nx")

    assert not os.path.exists(output_file), "output_file exists already"
    configuration = H52nxModel()

    configuration.output_file = output_file
    configuration.input_file = input_file
    configuration.raises_error = False
    configuration.no_master_file = False
    configuration.translation_y_keys = _cast_keys_section_to_tuple(("sy+spy",))
    configuration.translation_z_keys = _cast_keys_section_to_tuple(("sz+spz",))

    result = converter.from_h5_to_nx(configuration=configuration)
    assert os.path.exists(output_file)

    # check we have our 2 NXtomo are converted with (two distances).
    assert len(result) == 2
    # expected:
    nx_tomos = [NXtomo().load(file_path, entry) for (file_path, entry) in result]
    n_proj_per_nx_tomo = 1000
    n_frames = n_proj_per_nx_tomo + 20 * 3 + 3
    nx_tomo = nx_tomos[0]
    # check instrument
    assert len(nx_tomo.instrument.detector.image_key_control) == n_frames
    # image key
    numpy.testing.assert_array_equal(
        nx_tomo.instrument.detector.image_key_control,
        numpy.concatenate(
            (
                [ImageKey.DARK_FIELD] * 20,
                [ImageKey.FLAT_FIELD] * 20,
                [ImageKey.PROJECTION] * n_proj_per_nx_tomo,
                [ImageKey.ALIGNMENT] * 3,
                [ImageKey.FLAT_FIELD] * 20,
            )
        ),
    )

    # sequence number
    numpy.testing.assert_array_equal(
        nx_tomo.instrument.detector.sequence_number, numpy.arange(0, n_frames)
    )

    # detector pixel size
    numpy.testing.assert_almost_equal(
        nx_tomo.instrument.detector.x_pixel_size.magnitude,
        (2.95203 * _ureg.micrometer).magnitude,
        decimal=5,
    )
    numpy.testing.assert_almost_equal(
        nx_tomo.instrument.detector.y_pixel_size.magnitude,
        (2.95203 * _ureg.micrometer).magnitude,
        decimal=5,
    )

    assert numpy.isclose(
        nx_tomo.instrument.detector.distance, 1206.69 * _ureg.millimeter
    )
    # check energy
    assert nx_tomo.energy == 17.1 * _ureg.keV
    # check sample
    numpy.testing.assert_almost_equal(
        nx_tomo.sample.x_pixel_size.magnitude,
        (0.025 * _ureg.micrometer).magnitude,
        decimal=5,
    )
    numpy.testing.assert_almost_equal(
        nx_tomo.sample.y_pixel_size.magnitude,
        (0.025 * _ureg.micrometer).magnitude,
        decimal=5,
    )

    # check source
    numpy.testing.assert_almost_equal(
        nx_tomo.instrument.source.distance.magnitude,
        (-10.3065 * _ureg.millimeter).magnitude,
        decimal=4,
    )

    with h5py.File(input_file, mode="r") as h5f:
        scan_prj_index = 5
        sx = h5f[f"{scan_prj_index}.1/instrument/positioners/sx"][()] * _ureg.millimeter
        sy = h5f[f"{scan_prj_index}.1/instrument/positioners/sy"][()] * _ureg.millimeter
        spy = (
            h5f[f"{scan_prj_index}.1/instrument/positioners/spy"][()] * _ureg.micrometer
        )
        sz = h5f[f"{scan_prj_index}.1/instrument/positioners/sz"][()] * _ureg.millimeter
        spz = (
            h5f[f"{scan_prj_index}.1/instrument/positioners/spz"][()] * _ureg.micrometer
        )
        sax = (
            h5f[f"{scan_prj_index}.1/instrument/positioners/sax"][()] * _ureg.millimeter
        )
        say = (
            h5f[f"{scan_prj_index}.1/instrument/positioners/say"][()] * _ureg.millimeter
        )

    proj_index = numpy.arange(40, n_proj_per_nx_tomo + 40)

    numpy.testing.assert_almost_equal(
        nx_tomo.sample.x_translation.to_base_units().magnitude[proj_index],
        (sy + spy + say).to_base_units().magnitude,
    )
    numpy.testing.assert_almost_equal(
        nx_tomo.sample.y_translation.to_base_units().magnitude[proj_index],
        (sz + spz).to_base_units().magnitude,
    )
    numpy.testing.assert_almost_equal(
        (nx_tomo.sample.z_translation).to_base_units().magnitude[proj_index],
        (sx + sax).to_base_units().magnitude,
    )


def test_holotomo_conversion_id16b(tmp_path):
    """
    Test conversion of files from
    """
    dataset_file = os.path.join(
        "h5_datasets", "holotomo", "id16b", "20250711", "B3P7_sponge_pco2_ht_25nm.h5"
    )
    input_file = GitlabDataset.get_dataset(dataset_file)
    assert os.path.exists(input_file)
    output_file = os.path.join(tmp_path, "nexus_file.nx")

    configuration = H52nxModel()

    configuration.output_file = output_file
    configuration.input_file = input_file
    configuration.raises_error = False
    configuration.no_master_file = False

    result = converter.from_h5_to_nx(configuration=configuration)
    assert os.path.exists(output_file)

    # check we have our 4 NXtomo are converted with (two distances).
    assert len(result) == 4
    # expected:
    nx_tomos = [NXtomo().load(file_path, entry) for (file_path, entry) in result]
    n_proj_per_nx_tomo = 3203
    n_frames = n_proj_per_nx_tomo + (50 + 51) * 2 + 5

    nx_tomo = nx_tomos[0]
    # check instrument
    assert len(nx_tomo.instrument.detector.image_key_control) == n_frames
    # image key
    numpy.testing.assert_array_equal(
        nx_tomo.instrument.detector.image_key_control,
        numpy.concatenate(
            (
                [ImageKey.DARK_FIELD] * 50,
                [ImageKey.FLAT_FIELD] * 51,
                [ImageKey.PROJECTION] * n_proj_per_nx_tomo,
                [ImageKey.ALIGNMENT] * 5,
                [ImageKey.FLAT_FIELD] * 51,
                [ImageKey.DARK_FIELD] * 50,
            )
        ),
    )

    # sequence number
    numpy.testing.assert_array_equal(
        nx_tomo.instrument.detector.sequence_number, numpy.arange(0, n_frames)
    )

    # detector pixel size
    numpy.testing.assert_almost_equal(
        nx_tomo.instrument.detector.x_pixel_size.magnitude,
        (0.65 * _ureg.micrometer).magnitude,
        decimal=5,
    )
    numpy.testing.assert_almost_equal(
        nx_tomo.instrument.detector.y_pixel_size.magnitude,
        (0.65 * _ureg.micrometer).magnitude,
        decimal=5,
    )

    assert numpy.isclose(
        nx_tomo.instrument.detector.distance, 723.908 * _ureg.millimeter
    )
    # check energy
    assert nx_tomo.energy == 29.6 * _ureg.keV
    # check sample
    numpy.testing.assert_almost_equal(
        nx_tomo.sample.x_pixel_size.magnitude,
        (0.0250011 * _ureg.micrometer).magnitude,
        decimal=5,
    )
    numpy.testing.assert_almost_equal(
        nx_tomo.sample.y_pixel_size.magnitude,
        (0.0250011 * _ureg.micrometer).magnitude,
        decimal=5,
    )

    # check source
    numpy.testing.assert_almost_equal(
        nx_tomo.instrument.source.distance.magnitude,
        (-28.9577 * _ureg.millimeter).magnitude,
        decimal=4,
    )

    with h5py.File(input_file, mode="r") as h5f:
        scan_prj_index = 4
        sx = h5f[f"{scan_prj_index}.1/instrument/positioners/sx"][()] * _ureg.millimeter
        sy = h5f[f"{scan_prj_index}.1/instrument/positioners/sy"][()] * _ureg.millimeter
        sz = h5f[f"{scan_prj_index}.1/instrument/positioners/sz"][()] * _ureg.millimeter

    proj_index = numpy.arange(101, n_proj_per_nx_tomo + 101)

    numpy.testing.assert_almost_equal(
        nx_tomo.sample.x_translation.to_base_units().magnitude[proj_index],
        sy.to_base_units().magnitude,
    )
    numpy.testing.assert_almost_equal(
        nx_tomo.sample.y_translation.to_base_units().magnitude[proj_index],
        sz.to_base_units().magnitude,
    )
    numpy.testing.assert_almost_equal(
        (nx_tomo.sample.z_translation).to_base_units().magnitude[proj_index],
        sx.to_base_units().magnitude,
    )
