from __future__ import annotations

import os
import pint
import numpy
from datetime import datetime, timezone

from nxtomomill.tests.datasets import GitlabDataset
from nxtomomill import converter
from nxtomomill.io.config import TomoHDF5Config
from nxtomo import NXtomo

_ureg = pint.get_application_registry()


def test_multitomo_z_series(tmp_path):
    """test nxtomomill h52nx with a 'multitomo z-series' acquisition"""

    scan_dir = GitlabDataset.get_dataset("h5_datasets")
    assert scan_dir is not None
    input_file = os.path.join(
        scan_dir, "multitomo/Zseries_Volcano_passedevis_-1.7-aka720-bis_0002.h5"
    )

    configuration = TomoHDF5Config()

    configuration.input_file = input_file
    configuration.output_file = os.path.join(tmp_path, "nexus_file.nx")

    result = converter.from_h5_to_nx(configuration=configuration)
    assert len(result) == 20

    nx_tomo = NXtomo().load(result[0][0], result[0][1])
    assert len(nx_tomo.instrument.detector.image_key) == 3600
    desired_angles = (
        numpy.concatenate(
            (
                numpy.zeros(101),
                numpy.linspace(0, -360, 3499, endpoint=False),
            )
        )
        * _ureg.degree
    )
    numpy.testing.assert_almost_equal(
        nx_tomo.sample.rotation_angle.to_base_units().magnitude,
        desired_angles.to_base_units().magnitude,
        decimal=3,
    )

    # check start_time and end_time
    nx_tomos = [NXtomo().load(item[0], item[1]) for item in result]
    estimated_start_times = (
        datetime.fromisoformat("2026-02-24T15:47:23.217000").replace(
            tzinfo=timezone.utc
        ),
        datetime.fromisoformat("2026-02-24T15:47:55.729000").replace(
            tzinfo=timezone.utc
        ),
        datetime.fromisoformat("2026-02-24T15:48:28.250000").replace(
            tzinfo=timezone.utc
        ),
        datetime.fromisoformat("2026-02-24T15:49:00.771000").replace(
            tzinfo=timezone.utc
        ),
        datetime.fromisoformat("2026-02-24T15:49:33.292000").replace(
            tzinfo=timezone.utc
        ),
        datetime.fromisoformat("2026-02-24T15:50:05.813000").replace(
            tzinfo=timezone.utc
        ),
        datetime.fromisoformat("2026-02-24T15:50:38.334000").replace(
            tzinfo=timezone.utc
        ),
        datetime.fromisoformat("2026-02-24T15:51:10.855000").replace(
            tzinfo=timezone.utc
        ),
    )
    estimated_stop_times = (
        datetime.fromisoformat("2026-02-24T15:47:55.719000").replace(
            tzinfo=timezone.utc
        ),
        datetime.fromisoformat("2026-02-24T15:48:28.240000").replace(
            tzinfo=timezone.utc
        ),
        datetime.fromisoformat("2026-02-24T15:49:00.761000").replace(
            tzinfo=timezone.utc
        ),
        datetime.fromisoformat("2026-02-24T15:49:33.282000").replace(
            tzinfo=timezone.utc
        ),
        datetime.fromisoformat("2026-02-24T15:50:05.803000").replace(
            tzinfo=timezone.utc
        ),
        datetime.fromisoformat("2026-02-24T15:50:38.324000").replace(
            tzinfo=timezone.utc
        ),
        datetime.fromisoformat("2026-02-24T15:51:10.845000").replace(
            tzinfo=timezone.utc
        ),
        datetime.fromisoformat("2026-02-24T15:51:43.366000").replace(
            tzinfo=timezone.utc
        ),
    )

    for start_time, end_time, nx_tomo in zip(
        estimated_start_times, estimated_stop_times, nx_tomos
    ):
        assert start_time == nx_tomo.start_time
        assert end_time == nx_tomo.end_time
