import os
import tempfile

import h5py
import numpy.random
import pytest
from silx.io.url import DataUrl
from nxtomo.io import cwd_context
from pyunitsystem.voltagesystem import VoltageSystem
from nxtomo.utils.transformation import TransformationAxis

from nxtomo.nxobject.nxtransformations import Transformation
from nxtomo.nxobject.nxdetector import (
    FieldOfView,
    NXdetector,
    NXdetectorWithUnit,
    ImageKey,
)


def test_nx_detector():
    """test creation and saving of an nxdetector"""
    nx_detector = NXdetector(expected_dim=(2, 3))

    # check data
    with pytest.raises(TypeError):
        nx_detector.data = 12
    # if expected dims is not fulfill
    with pytest.raises(ValueError):
        nx_detector.data = numpy.random.random(100 * 100 * 5).reshape(5, 10, 10, 100)
    with pytest.raises(TypeError):
        nx_detector.data = (
            12,
            13,
        )
    nx_detector.data = numpy.random.random(100 * 100 * 5).reshape(5, 100, 100)

    # check image key control
    with pytest.raises(TypeError):
        nx_detector.image_key_control = 12
    nx_detector.image_key_control = [1] * 5
    nx_detector.image_key_control = [ImageKey.PROJECTION] * 5

    # check x and y pixel size (both 'real' and 'sample')
    with pytest.raises(TypeError):
        nx_detector.x_pixel_size = "test"
    nx_detector.x_pixel_size = 1e-7

    with pytest.raises(TypeError):
        nx_detector.y_pixel_size = {}
    nx_detector.y_pixel_size = 2e-7

    # check x_flipped and y_flipped
    with pytest.raises(TypeError):
        nx_detector.x_flipped = 12
    with pytest.raises(TypeError):
        nx_detector.y_flipped = 12

    nx_detector.x_flipped = True
    nx_detector.y_flipped = False

    # check detector distance
    with pytest.raises(TypeError):
        nx_detector.distance = "test"
    nx_detector.distance = 0.02

    # check field of view
    with pytest.raises(ValueError):
        nx_detector.field_of_view = "test"
    nx_detector.field_of_view = FieldOfView.HALF

    # check count time
    with pytest.raises(TypeError):
        nx_detector.count_time = 12
    nx_detector.count_time = [0.1] * 5

    # check estimated cor from motor
    with pytest.raises(TypeError):
        nx_detector.estimated_cor_from_motor = "test"
    nx_detector.estimated_cor_from_motor = 0.5

    assert isinstance(nx_detector.to_nx_dict(), dict)

    # check we can't set undefined attributes
    with pytest.raises(AttributeError):
        nx_detector.test = 12

    # test nx_detector concatenation
    concatenated_nx_detector = NXdetector.concatenate([nx_detector, nx_detector])
    numpy.testing.assert_array_equal(
        concatenated_nx_detector.image_key_control, [ImageKey.PROJECTION] * 10
    )
    assert concatenated_nx_detector.x_pixel_size.value == 1e-7
    assert concatenated_nx_detector.y_pixel_size.value == 2e-7
    assert concatenated_nx_detector.distance.value == 0.02
    nx_detector.field_of_view = FieldOfView.HALF
    nx_detector.count_time = [0.1] * 10
    nx_detector.estimated_cor_from_motor = 0.5


def test_nx_detector_with_unit():
    diode = NXdetectorWithUnit(
        node_name="diode",
        expected_dim=(1,),
        default_unit=VoltageSystem.VOLT,
    )
    with pytest.raises(ValueError):
        diode.data = numpy.arange(10 * 10).reshape([10, 10])
    with pytest.raises(TypeError):
        diode.data = [10, 12]
    with pytest.raises(TypeError):
        diode.data = "test"
    diode.data = None
    diode.data = numpy.random.random(12)
    diode.data = (DataUrl(),)

    # test nx_detector concatenation
    concatenated_nx_detector = NXdetectorWithUnit.concatenate(
        [diode, diode],
        expected_dim=(1,),
        default_unit=VoltageSystem.VOLT,
    )
    assert len(concatenated_nx_detector.data.value) == 2
    assert isinstance(concatenated_nx_detector.data.value[1], DataUrl)


def test_nx_detector_with_virtual_source():
    """Insure detector data can be write from Virtual sources"""
    cwd = os.getcwd()
    with tempfile.TemporaryDirectory() as tmp_folder:
        # create virtual dataset
        n_base_raw_dataset = 5
        n_z, n_y, n_x = 4, 100, 100
        base_raw_dataset_shape = (n_z, n_y, n_x)
        n_base_raw_dataset_elmts = n_z * n_y * n_x

        v_sources = []

        raw_files = [
            os.path.join(tmp_folder, f"raw_file_{i_file}.hdf5")
            for i_file in range(n_base_raw_dataset)
        ]
        for i_raw_file, raw_file in enumerate(raw_files):
            with h5py.File(raw_file, mode="w") as h5f:
                h5f["data"] = numpy.arange(
                    start=n_base_raw_dataset_elmts * i_raw_file,
                    stop=n_base_raw_dataset_elmts * (i_raw_file + 1),
                ).reshape(base_raw_dataset_shape)
                v_sources.append(h5py.VirtualSource(h5f["data"]))

        nx_detector = NXdetector()
        nx_detector.data = v_sources

        detector_file = os.path.join(tmp_folder, "detector_file.hdf5")
        nx_detector.save(file_path=detector_file)

        # check the virtual dataset has been properly createde and linked
        with h5py.File(detector_file, mode="r") as h5f_master:
            dataset = h5f_master["/detector/data"]
            assert dataset.is_virtual
            for i_raw_file, raw_file in enumerate(raw_files):
                with h5py.File(raw_file, mode="r") as h5f_raw:
                    numpy.testing.assert_array_equal(
                        dataset[i_raw_file * n_z : (i_raw_file + 1) * n_z],
                        h5f_raw["data"],
                    )
            # check attributes have beem rewrite as expected
            assert "interpretation" in dataset.attrs

            # check virtual dataset is composed of relative links
            for vs_info in dataset.virtual_sources():
                assert vs_info.file_name.startswith("./")
        assert cwd == os.getcwd()

    # check concatenation
    concatenated_nx_detector = NXdetector.concatenate([nx_detector, nx_detector])
    assert isinstance(concatenated_nx_detector.data[1], h5py.VirtualSource)
    assert len(concatenated_nx_detector.data) == len(raw_files) * 2


def test_nx_detector_with_local_urls():
    """Insure detector data can be write from DataUrl linking to local dataset (in the same file)"""

    cwd = os.getcwd()
    n_base_dataset = 3
    n_z, n_y, n_x = 2, 10, 20
    base_dataset_shape = (n_z, n_y, n_x)
    n_base_dataset_elmts = n_z * n_y * n_x
    urls = []

    with tempfile.TemporaryDirectory() as tmp_folder:
        master_file = os.path.join(tmp_folder, "master_file.hdf5")
        with h5py.File(master_file, mode="a") as h5f:
            for i in range(n_base_dataset):
                data_path = f"/data_{i}"
                h5f[data_path] = numpy.arange(
                    start=n_base_dataset_elmts * i,
                    stop=n_base_dataset_elmts * (i + 1),
                ).reshape(base_dataset_shape)
                urls.append(
                    DataUrl(
                        file_path=master_file,
                        data_path=data_path,
                        scheme="silx",
                    )
                )
        nx_detector = NXdetector()
        nx_detector.data = urls
        nx_detector.save(file_path=master_file)

        # check the virtual dataset has been properly createde and linked
        with h5py.File(master_file, mode="r") as h5f_master:
            dataset = h5f_master["/detector/data"]
            assert dataset.is_virtual
            for i in range(n_base_dataset):
                numpy.testing.assert_array_equal(
                    dataset[i * n_z : (i + 1) * n_z],
                    numpy.arange(
                        start=n_base_dataset_elmts * i,
                        stop=n_base_dataset_elmts * (i + 1),
                    ).reshape(base_dataset_shape),
                )
            # check virtual dataset is composed of relative links
            for vs_info in dataset.virtual_sources():
                assert vs_info.file_name.startswith("./")
        assert cwd == os.getcwd()

    # check concatenation
    concatenated_nx_detector = NXdetector.concatenate([nx_detector, nx_detector])
    assert isinstance(concatenated_nx_detector.data[1], DataUrl)
    assert len(concatenated_nx_detector.data) == n_base_dataset * 2


def test_nx_detector_with_external_urls():
    """Insure detector data can be write from DataUrl linking to external dataset"""
    cwd = os.getcwd()
    with tempfile.TemporaryDirectory() as tmp_folder:
        # create virtual dataset
        n_base_raw_dataset = 5
        n_z, n_y, n_x = 4, 100, 100
        base_raw_dataset_shape = (n_z, n_y, n_x)
        n_base_raw_dataset_elmts = n_z * n_y * n_x

        urls = []

        raw_files = [
            os.path.join(tmp_folder, f"raw_file_{i_file}.hdf5")
            for i_file in range(n_base_raw_dataset)
        ]
        for i_raw_file, raw_file in enumerate(raw_files):
            with h5py.File(raw_file, mode="w") as h5f:
                h5f["data"] = numpy.arange(
                    start=n_base_raw_dataset_elmts * i_raw_file,
                    stop=n_base_raw_dataset_elmts * (i_raw_file + 1),
                ).reshape(base_raw_dataset_shape)
                # provide one file path each two as an absolue path
                if i_raw_file % 2 == 0:
                    file_path = os.path.abspath(raw_file)
                else:
                    file_path = os.path.relpath(raw_file, tmp_folder)
                urls.append(
                    DataUrl(
                        file_path=file_path,
                        data_path="data",
                        scheme="silx",
                    )
                )

        nx_detector = NXdetector()
        nx_detector.data = urls

        detector_file = os.path.join(tmp_folder, "detector_file.hdf5")
        # needed as we provide some link with relative path
        with cwd_context(tmp_folder):
            nx_detector.save(file_path=detector_file)

        # check the virtual dataset has been properly createde and linked
        with h5py.File(detector_file, mode="r") as h5f_master:
            dataset = h5f_master["/detector/data"]
            assert dataset.is_virtual
            for i_raw_file, raw_file in enumerate(raw_files):
                with h5py.File(raw_file, mode="r") as h5f_raw:
                    numpy.testing.assert_array_equal(
                        dataset[i_raw_file * n_z : (i_raw_file + 1) * n_z],
                        h5f_raw["data"],
                    )
            # check virtual dataset is composed of relative links
            for vs_info in dataset.virtual_sources():
                assert vs_info.file_name.startswith("./")

        assert cwd == os.getcwd()
        # check concatenation
        concatenated_nx_detector = NXdetector.concatenate([nx_detector, nx_detector])
        assert isinstance(concatenated_nx_detector.data[1], DataUrl)
        assert len(concatenated_nx_detector.data) == n_base_raw_dataset * 2


@pytest.mark.parametrize(
    "load_data_as, expected_type",
    [
        ("as_numpy_array", numpy.ndarray),
        ("as_virtual_source", h5py.VirtualSource),
        ("as_data_url", DataUrl),
    ],
)
def test_load_detector_data(tmp_path, load_data_as, expected_type):
    layout = h5py.VirtualLayout(shape=(4 * 2, 100, 100), dtype="i4")

    for n in range(0, 4):
        filename = os.path.join(tmp_path, "{n}.h5")
        with h5py.File(filename, "w") as f:
            f["data"] = numpy.arange(100 * 100 * 2).reshape(2, 100, 100)

        vsource = h5py.VirtualSource(filename, "data", shape=(2, 100, 100))
        start_n = n * 2
        end_n = start_n + 2
        layout[start_n:end_n] = vsource

    output_file = os.path.join(tmp_path, "VDS.h5")
    with h5py.File(output_file, "w") as f:
        f.create_virtual_dataset("data", layout, fillvalue=-5)

    nx_detector = NXdetector()
    nx_detector._load(
        file_path=output_file,
        data_path="/",
        load_data_as=load_data_as,
        nexus_version=None,
    )

    if expected_type is numpy.ndarray:
        assert isinstance(nx_detector.data, expected_type)
    else:
        for elmt in nx_detector.data:
            assert isinstance(elmt, expected_type)
    nx_detector.save(os.path.join(tmp_path, "output_file.nx"))


def test_nxtransformations_with_nxdetector(tmp_path):
    """
    test behavior of an Nxtransformations with an NXtomo
    and coherence between lr_flip / ud_flip API (provided for convenience) and providing directly the
    transformations
    """

    def build_detector():
        nx_detector = NXdetector(expected_dim=(2, 3))
        nx_detector.data = numpy.random.random(100 * 100 * 5).reshape(5, 100, 100)
        nx_detector.image_key_control = [1] * 5
        nx_detector.image_key_control = [ImageKey.PROJECTION] * 5
        return nx_detector

    nx_detector_1 = build_detector()
    nx_detector_2 = build_detector()

    # test having a left-right flip
    nx_detector_1.transformations.addTransformation(
        Transformation(
            axis_name="rz",
            value=180,
            transformation_type="rotation",
            vector=TransformationAxis.AXIS_Z,
        )
    )
    nx_detector_2.x_flipped = True

    assert (
        nx_detector_1.transformations.to_nx_dict()
        == nx_detector_2.transformations.to_nx_dict()
    )

    # test having a up-down flip
    nx_detector_3 = build_detector()
    nx_detector_4 = build_detector()
    nx_detector_3.transformations.addTransformation(
        Transformation(
            axis_name="ry",
            value=180,
            transformation_type="rotation",
            vector=TransformationAxis.AXIS_Y,
        )
    )
    nx_detector_4.y_flipped = True

    assert (
        nx_detector_3.transformations.to_nx_dict()
        == nx_detector_4.transformations.to_nx_dict()
    )

    # having both lr and ud
    nx_detector_5 = build_detector()
    nx_detector_6 = build_detector()

    nx_detector_5.transformations.addTransformation(
        Transformation(
            axis_name="ry",
            value=180,
            transformation_type="rotation",
            vector=TransformationAxis.AXIS_Y,
        )
    )
    nx_detector_5.transformations.addTransformation(
        Transformation(
            axis_name="rz",
            value=180,
            transformation_type="rotation",
            vector=TransformationAxis.AXIS_Z,
        )
    )
    nx_detector_6.x_flipped = True
    nx_detector_6.y_flipped = True

    assert (
        nx_detector_5.transformations.to_nx_dict()
        == nx_detector_6.transformations.to_nx_dict()
    )
