from __future__ import annotations

import pint
import pytest
import h5py
import numpy

from nxtomomill.converter.hdf5.utils import (
    get_default_output_file,
    PROCESSED_DATA_DIR_NAME,
    RAW_DATA_DIR_NAME,
)
from nxtomomill.models.h52nx.h52nx import H52nxModel
from nxtomomill.converter.hdf5.acquisition.utils._context import (
    Proxy,
    _BeamlineContext,
    update_config_keys_from_beamline_context,
    update_config_keys_from_tomo_config_context,
    read_key,
    MissingValue,
)
from nxtomomill.converter.hdf5.acquisition.blisstomoconfig import TomoConfig
from nxtomomill.converter.hdf5.acquisition.utils._keys import _Key

_ureg = pint.UnitRegistry()


def test_get_default_output_file():
    """
    test the get_default_output_file function
    """
    # 1. simple test
    assert get_default_output_file("/my_tmp/path/file.h5") == "/my_tmp/path/file.nx"
    assert (
        get_default_output_file(f"/my_tmp/{RAW_DATA_DIR_NAME}/file.h5")
        == f"/my_tmp/{PROCESSED_DATA_DIR_NAME}/file.nx"
    )

    assert (
        get_default_output_file(f"/my_tmp/path/{RAW_DATA_DIR_NAME}/toto/file.h5")
        == f"/my_tmp/path/{PROCESSED_DATA_DIR_NAME}/toto/file.nx"
    )
    # note: _RAW_DATA_DIR_NAME part of the path but not a folder
    assert (
        get_default_output_file(f"/my_tmp/path_{RAW_DATA_DIR_NAME}/toto/file.h5")
        == f"/my_tmp/path_{RAW_DATA_DIR_NAME}/toto/file.nx"
    )

    # 2. advance test
    # 2.1 use case: '_RAW_DATA_DIR_NAME' is present twice in the path -> replace the deeper one
    assert (
        get_default_output_file(
            f"/my_tmp/{RAW_DATA_DIR_NAME}/path/{RAW_DATA_DIR_NAME}/toto/file.h5"
        )
        == f"/my_tmp/{RAW_DATA_DIR_NAME}/path/{PROCESSED_DATA_DIR_NAME}/toto/file.nx"
    )

    # 2.2 use case: contains both '_RAW_DATA_DIR_NAME' and '_PROCESSED_DATA_DIR_NAME' in the path
    assert (
        get_default_output_file(
            f"/my_tmp/{RAW_DATA_DIR_NAME}/path/{PROCESSED_DATA_DIR_NAME}/toto/file.h5"
        )
        == f"/my_tmp/{RAW_DATA_DIR_NAME}/path/{PROCESSED_DATA_DIR_NAME}/toto/file.nx"
    )

    assert (
        get_default_output_file(
            f"/my_tmp/{PROCESSED_DATA_DIR_NAME}/path/{RAW_DATA_DIR_NAME}/toto/file.h5"
        )
        == f"/my_tmp/{PROCESSED_DATA_DIR_NAME}/path/{PROCESSED_DATA_DIR_NAME}/toto/file.nx"
    )

    # 2.3 use case: expected output file is the input file. Make sure append '_nxtomo'
    assert (
        get_default_output_file("/my_tmp/path/file.nx") == "/my_tmp/path/file_nxtomo.nx"
    )


def test_update_config_keys_from_tomo_config_context():
    tomo_config = TomoConfig()
    # tomoconfig register motors as (real_motor_name, bliss_alias)
    tomo_config.sample_x = ("alias_sx", "sx")
    tomo_config.sample_y = ("alias_sy", "sy")
    tomo_config.translation_y = ("alias_ty", "ty")

    model = H52nxModel()
    model.sample_x_keys = (_Key(dataset="sample_x", contexts={Proxy.TOMO_CONFIG}),)
    model.sample_y_keys = (
        _Key(dataset="sample_y+translation_y", contexts={Proxy.TOMO_CONFIG}),
    )

    update_config_keys_from_tomo_config_context(model, tomo_config=tomo_config)
    assert model.sample_x_keys == (_Key(dataset="sx", contexts=set()),)
    assert model.sample_y_keys == (_Key(dataset="sy+ty", contexts=set()),)


def test_update_config_keys_from_beamline_context():
    model = H52nxModel()
    model.sample_x_keys = ("esrf-id11:sx", "ud", "id03:tx")

    update_config_keys_from_beamline_context(model, _BeamlineContext.id11)
    assert model.sample_x_keys == (
        _Key(dataset="sx", contexts={_BeamlineContext.id11}),
        _Key(dataset="ud", contexts=set()),
    )

    update_config_keys_from_beamline_context(model, None)
    # if no beamline found we want to get rid of all the keys with a beamline context
    assert model.sample_x_keys == (_Key(dataset="ud", contexts=set()),)


@pytest.fixture
def read_key_datasets(tmp_path):
    """
    Create an HDF5 file with datasets of coherent and incoherent sizes,
    some with units and some without.
    """
    # Create a temporary HDF5 file
    file_path = tmp_path / "test_datasets.h5"

    with h5py.File(file_path, "w") as f:
        # Create datasets with coherent sizes (same shape)
        f.create_dataset("coherent/dataset1", data=numpy.random.rand(10))
        f.create_dataset("coherent/dataset2", data=numpy.random.rand(10))

        # Create datasets with units (using pint)
        f.create_dataset("with_incoherent_units/dataset5", data=numpy.random.rand(10))
        f["with_incoherent_units/dataset5"].attrs["unit"] = "meter"

        f.create_dataset("with_incoherent_units/dataset6", data=numpy.random.rand(10))
        f["with_incoherent_units/dataset6"].attrs["unit"] = "second"

        # Create datasets without units
        f.create_dataset("without_units/dataset7", data=numpy.ones(10))
        f.create_dataset("without_units/dataset8", data=numpy.ones(10))

        # Create datasets with **DIS**
        f.create_dataset("missing_data/dataset9", data="*DIS*")
        f.create_dataset("missing_data/dataset10", data=numpy.random.random(10))

        # Create datasets with coherent units (but different)
        f.create_dataset("with_coherent_units/dataset11", data=numpy.ones(10))
        f["with_coherent_units/dataset11"].attrs["unit"] = "meter"

        f.create_dataset("with_coherent_units/dataset12", data=(numpy.ones(10) * 2))
        f["with_coherent_units/dataset12"].attrs["unit"] = "cm"
        f.create_dataset("with_coherent_units/dataset13", data=numpy.ones(10))
        f["with_coherent_units/dataset13"].attrs["unit"] = "cm"

        # Create datasets with coherent units (arrays with different size)
        f.create_dataset("with_different_nb_points/dataset14", data=numpy.ones(10))
        f["with_different_nb_points/dataset14"].attrs["unit"] = "meter"

        f.create_dataset("with_different_nb_points/dataset15", data=(numpy.ones(6) * 2))
        f["with_different_nb_points/dataset15"].attrs["unit"] = "cm"

        # Create datasets with coherent units (arrays and scalars)
        f.create_dataset("with_scalar_and_arrays/dataset16", data=numpy.ones(10))
        f["with_scalar_and_arrays/dataset16"].attrs["unit"] = "meter"

        f.create_dataset("with_scalar_and_arrays/dataset17", data=1.4)
        f["with_scalar_and_arrays/dataset17"].attrs["unit"] = "cm"

        f.create_dataset("with_scalar_and_arrays/dataset18", data=(numpy.ones(10) * 2))
        f["with_scalar_and_arrays/dataset18"].attrs["unit"] = "meter"

        # Create datasets with coherent units (arrays and scalars)
        f.create_dataset(
            "with_scalar_and_incoherent_arrays/dataset19", data=numpy.ones(10)
        )
        f["with_scalar_and_incoherent_arrays/dataset19"].attrs["unit"] = "meter"

        f.create_dataset("with_scalar_and_incoherent_arrays/dataset20", data=1.4)
        f["with_scalar_and_incoherent_arrays/dataset20"].attrs["unit"] = "cm"

        f.create_dataset(
            "with_scalar_and_incoherent_arrays/dataset21", data=numpy.ones(9)
        )
        f["with_scalar_and_incoherent_arrays/dataset21"].attrs["unit"] = "meter"

    return file_path


def test_read_key(read_key_datasets, caplog):
    with h5py.File(read_key_datasets) as h5f:
        # test with dataset not existing
        with pytest.raises(KeyError):
            read_key(
                key="dataset3+dataset4",
                node=h5f["coherent"],
                info_retrieve="",
                expected_unit=_ureg.meter,
            )

        # test with dataset not dump
        with pytest.raises(MissingValue):
            read_key(
                key="dataset9+dataset10",
                node=h5f["missing_data"],
                info_retrieve="",
                expected_unit=_ureg.meter,
            )

        # test with dataset with coherent size and no unit
        quantity: pint.Quantity = read_key(
            key="dataset7+dataset8",
            node=h5f["without_units"],
            info_retrieve="",
            expected_unit=_ureg.centimeter,
        )
        numpy.testing.assert_array_equal(
            quantity.to_base_units(), numpy.array([0.02] * 10)
        )

        # test with dataset with coherent size but incoherent units
        with pytest.raises(pint.errors.DimensionalityError):
            quantity: pint.Quantity = read_key(
                key="dataset5+dataset6",
                node=h5f["with_incoherent_units"],
                info_retrieve="",
                expected_unit=_ureg.centimeter,
            )

        # test with dataset with coherent size and coherent units
        quantity: pint.Quantity = read_key(
            key="dataset11+dataset12-dataset13",
            node=h5f["with_coherent_units"],
            info_retrieve="",
            expected_unit=_ureg.centimeter,
        )
        numpy.testing.assert_array_equal(
            quantity.to_base_units(), numpy.array([1.01] * 10)
        )

        # test arrays with incoherent size
        key = "dataset14+dataset15"
        expected_warning = (
            "nxtomomill.converter.hdf5.acquisition.utils._context",
            30,
            f"Incoherent number of points founds for info when accessing '{key}'.",
        )
        assert expected_warning not in caplog.record_tuples
        quantity: pint.Quantity = read_key(
            key=key,
            node=h5f["with_different_nb_points"],
            info_retrieve="info",
            expected_unit=_ureg.centimeter,
        )
        numpy.testing.assert_array_equal(
            quantity.to_base_units(), numpy.array([1.02] * 6)
        )
        assert expected_warning in caplog.record_tuples

        # test arrays with scalar
        quantity: pint.Quantity = read_key(
            key="dataset16+dataset17-dataset18",
            node=h5f["with_scalar_and_arrays"],
            info_retrieve="",
            expected_unit=_ureg.centimeter,
        )
        numpy.testing.assert_array_equal(
            quantity.to_base_units(), numpy.array([1 + 0.014 - 2] * 10)
        )
        # test incoherent arrays with scalar
        quantity: pint.Quantity = read_key(
            key="dataset19-dataset20+dataset21",
            node=h5f["with_scalar_and_incoherent_arrays"],
            info_retrieve="",
            expected_unit=_ureg.centimeter,
        )
        numpy.testing.assert_array_equal(
            quantity.to_base_units(), numpy.array([1 - 0.014 + 1] * 9)
        )
