import os
from datetime import datetime
from tempfile import TemporaryDirectory
import numpy

import pytest
from silx.io.url import DataUrl

from nxtomo.application.nxtomo import NXtomo
from nxtomo.nxobject.nxdetector import FieldOfView, ImageKey
from nxtomo.nxobject.utils import concatenate


nexus_path_versions = (1.1, 1.0, None)


@pytest.mark.parametrize("nexus_path_version", nexus_path_versions)
def test_nx_tomo(nexus_path_version):
    nx_tomo = NXtomo(node_name="")

    # check start time
    with pytest.raises(TypeError):
        nx_tomo.start_time = 12
    nx_tomo.start_time = datetime.now()

    # check end time
    with pytest.raises(TypeError):
        nx_tomo.end_time = 12
    nx_tomo.end_time = datetime(2022, 2, 27)

    # check sample
    with pytest.raises(TypeError):
        nx_tomo.sample = "tata"

    # check detector
    with pytest.raises(TypeError):
        nx_tomo.instrument.detector = "tata"

    # check energy
    with pytest.raises(TypeError):
        nx_tomo.energy = "tata"
    nx_tomo.energy = 12.3

    # check group size
    with pytest.raises(TypeError):
        nx_tomo.group_size = "tata"
    nx_tomo.group_size = 3

    # check title
    with pytest.raises(TypeError):
        nx_tomo.title = 12
    nx_tomo.title = "title"

    # check instrument
    with pytest.raises(TypeError):
        nx_tomo.instrument = "test"

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

    # create detector for test
    projections = numpy.random.random(100 * 100 * 8).reshape([8, 100, 100])
    flats_1 = numpy.random.random(100 * 100 * 2).reshape([2, 100, 100])
    darks = numpy.random.random(100 * 100 * 3).reshape([3, 100, 100])
    flats_2 = numpy.random.random(100 * 100 * 2).reshape([2, 100, 100])
    alignment = numpy.random.random(100 * 100 * 1).reshape([1, 100, 100])
    nx_tomo.instrument.detector.data = numpy.concatenate(
        [
            darks,
            flats_1,
            projections,
            flats_2,
            alignment,
        ]
    )

    nx_tomo.instrument.detector.image_key_control = numpy.concatenate(
        [
            [ImageKey.DARK_FIELD] * 3,
            [ImageKey.FLAT_FIELD] * 2,
            [ImageKey.PROJECTION] * 8,
            [ImageKey.FLAT_FIELD] * 2,
            [ImageKey.ALIGNMENT] * 1,
        ]
    )

    nx_tomo.instrument.detector.x_pixel_size = (
        nx_tomo.instrument.detector.y_pixel_size
    ) = 1e-7
    nx_tomo.instrument.detector.distance = 0.2
    nx_tomo.instrument.detector.field_of_view = FieldOfView.HALF
    nx_tomo.instrument.detector.count_time = numpy.concatenate(
        [
            [0.2] * 3,  # darks
            [0.1] * 2,  # flats 1
            [0.1] * 8,  # projections
            [0.1] * 2,  # flats 2
            [0.1] * 1,  # alignment
        ]
    )

    # create sample for test
    nx_tomo.sample.name = "my sample"
    nx_tomo.sample.rotation_angle = numpy.concatenate(
        [
            [0.0] * 3,  # darks
            [0.0] * 2,  # flats 1
            numpy.linspace(0, 180, num=8, endpoint=False),  # projections
            [180.0] * 2,  # flats 2
            [0.0],  # alignement
        ]
    )

    n_frames = 3 + 2 + 8 + 2 + 1
    nx_tomo.sample.x_translation = [0.6] * n_frames
    nx_tomo.sample.y_translation = [0.2] * n_frames
    nx_tomo.sample.z_translation = [0.1] * n_frames
    assert nx_tomo.is_root is True
    assert nx_tomo.instrument.is_root is False
    assert (
        nx_tomo.root_path
        == nx_tomo.instrument.root_path
        == nx_tomo.instrument.detector.root_path
    )

    NXtomo.check_consistency(nx_tomo=nx_tomo, raises_error=True)

    with TemporaryDirectory() as folder:
        file_path = os.path.join(folder, "nexus_file.hdf5")

        nx_tomo.save(
            file_path=file_path,
            data_path="entry",
            nexus_path_version=nexus_path_version,
        )
        assert os.path.exists(file_path)

        # insure we can read it back
        scan = NXtomo().load(file_path, data_path="entry")
        assert (
            len(
                tuple(
                    filter(
                        lambda image_key: image_key is ImageKey.FLAT_FIELD,
                        scan.instrument.detector.image_key_control,
                    )
                )
            )
            == 4
        )
        assert (
            len(
                tuple(
                    filter(
                        lambda image_key: image_key is ImageKey.DARK_FIELD,
                        scan.instrument.detector.image_key_control,
                    )
                )
            )
            == 3
        )
        assert (
            len(
                tuple(
                    filter(
                        lambda image_key: image_key is ImageKey.PROJECTION,
                        scan.instrument.detector.image_key_control,
                    )
                )
            )
            == 8
        )
        assert (
            len(
                tuple(
                    filter(
                        lambda image_key: image_key is ImageKey.ALIGNMENT,
                        scan.instrument.detector.image_key_control,
                    )
                )
            )
            == 1
        )
        assert scan.energy.value == 12.3
        assert scan.instrument.detector.x_pixel_size.value == 1e-7
        assert scan.instrument.detector.y_pixel_size.value == 1e-7
        assert scan.instrument.detector.distance.value == 0.2
        assert scan.instrument.detector.field_of_view == FieldOfView.HALF
        assert scan.sample.name == "my sample"
        assert (
            len(scan.sample.x_translation.value)
            == len(scan.sample.y_translation.value)
            == len(scan.sample.z_translation.value)
            == n_frames
        )
        assert scan.sample.x_translation.value[0] == 0.6
        assert scan.sample.y_translation.value[0] == 0.2
        assert scan.sample.z_translation.value[0] == 0.1
        if nexus_path_version != 1.0:
            assert scan.instrument.source.name is not None
            assert scan.instrument.source.type is not None

        # try to load it from the disk
        loaded_nx_tomo = NXtomo("test").load(file_path=file_path, data_path="entry")
        assert isinstance(loaded_nx_tomo, NXtomo)
        assert loaded_nx_tomo.energy.value == nx_tomo.energy.value
        assert loaded_nx_tomo.energy.unit == nx_tomo.energy.unit
        assert loaded_nx_tomo.start_time == nx_tomo.start_time
        assert loaded_nx_tomo.end_time == nx_tomo.end_time
        assert (
            loaded_nx_tomo.instrument.detector.x_pixel_size.value
            == nx_tomo.instrument.detector.x_pixel_size.value
        )
        assert (
            loaded_nx_tomo.instrument.detector.x_pixel_size.unit
            == nx_tomo.instrument.detector.x_pixel_size.unit
        )
        assert (
            loaded_nx_tomo.instrument.detector.y_pixel_size.value
            == nx_tomo.instrument.detector.y_pixel_size.value
        )
        assert (
            loaded_nx_tomo.instrument.detector.field_of_view
            == nx_tomo.instrument.detector.field_of_view
        )
        numpy.testing.assert_array_equal(
            loaded_nx_tomo.instrument.detector.count_time.value,
            nx_tomo.instrument.detector.count_time.value,
        )
        assert (
            loaded_nx_tomo.instrument.detector.count_time.unit
            == nx_tomo.instrument.detector.count_time.unit
        )
        assert (
            loaded_nx_tomo.instrument.detector.distance.value
            == nx_tomo.instrument.detector.distance.value
        )
        assert (
            loaded_nx_tomo.instrument.detector.distance.unit
            == nx_tomo.instrument.detector.distance.unit
        )
        assert (
            loaded_nx_tomo.instrument.detector.estimated_cor_from_motor
            == nx_tomo.instrument.detector.estimated_cor_from_motor
        )

        numpy.testing.assert_array_equal(
            loaded_nx_tomo.instrument.detector.image_key_control,
            nx_tomo.instrument.detector.image_key_control,
        )
        numpy.testing.assert_array_equal(
            loaded_nx_tomo.instrument.detector.image_key,
            nx_tomo.instrument.detector.image_key,
        )

        assert loaded_nx_tomo.sample.name == nx_tomo.sample.name
        assert loaded_nx_tomo.sample.rotation_angle is not None
        numpy.testing.assert_array_almost_equal(
            loaded_nx_tomo.sample.rotation_angle, nx_tomo.sample.rotation_angle
        )
        numpy.testing.assert_array_almost_equal(
            loaded_nx_tomo.sample.x_translation.value,
            nx_tomo.sample.x_translation.value,
        )
        numpy.testing.assert_array_almost_equal(
            loaded_nx_tomo.sample.y_translation.value,
            nx_tomo.sample.y_translation.value,
        )
        numpy.testing.assert_array_almost_equal(
            loaded_nx_tomo.sample.z_translation.value,
            nx_tomo.sample.z_translation.value,
        )

        loaded_nx_tomo = NXtomo("test").load(
            file_path=file_path, data_path="entry", detector_data_as="as_numpy_array"
        )
        numpy.testing.assert_array_almost_equal(
            loaded_nx_tomo.instrument.detector.data,
            nx_tomo.instrument.detector.data,
        )
        loaded_nx_tomo = NXtomo("test").load(
            file_path=file_path, data_path="entry", detector_data_as="as_data_url"
        )
        assert isinstance(loaded_nx_tomo.instrument.detector.data[0], DataUrl)
        with pytest.raises(ValueError):
            # check an error is raise because the dataset is not virtual
            loaded_nx_tomo = NXtomo("test").load(
                file_path=file_path,
                data_path="entry",
                detector_data_as="as_virtual_source",
            )

        # test concatenation
        nx_tomo_concat = concatenate([loaded_nx_tomo, None, loaded_nx_tomo])
        concat_file = os.path.join(folder, "concatenated_nexus_file.hdf5")

        nx_tomo_concat.save(
            file_path=concat_file,
            data_path="myentry",
            nexus_path_version=nexus_path_version,
        )
        loaded_concatenated_nx_tomo = NXtomo("test").load(
            file_path=concat_file,
            data_path="myentry",
            detector_data_as="as_virtual_source",
        )
        numpy.testing.assert_array_almost_equal(
            loaded_concatenated_nx_tomo.sample.rotation_angle,
            numpy.concatenate(
                [
                    nx_tomo.sample.rotation_angle,
                    nx_tomo.sample.rotation_angle,
                ]
            ),
        )
        numpy.testing.assert_array_almost_equal(
            loaded_concatenated_nx_tomo.sample.x_translation.value,
            numpy.concatenate(
                [
                    nx_tomo.sample.x_translation.value,
                    nx_tomo.sample.x_translation.value,
                ]
            ),
        )

        with pytest.raises(TypeError):
            concatenate([1, 2])


@pytest.mark.parametrize("nexus_path_version", nexus_path_versions)
def test_nx_tomo_subselection(nexus_path_version):
    """
    test sub_select_from_projection_angle_range
    """
    nx_tomo = NXtomo()
    nx_tomo.energy = 12.3
    shape = (12, 12)
    data_dark = numpy.ones(shape)
    data_flat = numpy.ones(shape) * 2.0
    data_projection = numpy.ones(shape) * 3.0
    str(nx_tomo)
    nx_tomo.instrument.detector.data = numpy.concatenate(
        (
            data_dark,
            data_dark,
            data_flat,
            data_projection,
            data_projection,
            data_projection,
            data_flat,
            data_projection,
            data_projection,
            data_projection,
            data_flat,
        )
    )
    nx_tomo.instrument.detector.image_key_control = numpy.array(
        (
            ImageKey.DARK_FIELD,
            ImageKey.DARK_FIELD,
            ImageKey.FLAT_FIELD,
            ImageKey.PROJECTION,
            ImageKey.PROJECTION,
            ImageKey.PROJECTION,
            ImageKey.FLAT_FIELD,
            ImageKey.PROJECTION,
            ImageKey.PROJECTION,
            ImageKey.PROJECTION,
            ImageKey.FLAT_FIELD,
        )
    )

    original_angles = numpy.array(
        (
            0,
            0,
            0,
            10,
            20.5,
            22.5,
            180,
            180,
            200,
            300.2,
            300.2,
        )
    )
    nx_tomo.sample.rotation_angle = original_angles

    nx_tomo_sub_1 = NXtomo.sub_select_from_angle_offset(
        nx_tomo=nx_tomo,
        start_angle_offset=10,
        angle_interval=12.5,
        shift_angles=False,
    )

    numpy.testing.assert_equal(
        nx_tomo_sub_1.instrument.detector.image_key_control,
        numpy.array(
            (
                ImageKey.DARK_FIELD,
                ImageKey.DARK_FIELD,
                ImageKey.FLAT_FIELD,
                ImageKey.INVALID,
                ImageKey.PROJECTION,
                ImageKey.PROJECTION,
                ImageKey.FLAT_FIELD,
                ImageKey.INVALID,
                ImageKey.INVALID,
                ImageKey.INVALID,
                ImageKey.FLAT_FIELD,
            )
        ),
    )
    numpy.testing.assert_equal(
        nx_tomo_sub_1.sample.rotation_angle,
        original_angles,
    )

    nx_tomo_sub_2 = NXtomo.sub_select_from_angle_offset(
        nx_tomo=nx_tomo,
        start_angle_offset=10,
        angle_interval=20,
        shift_angles=True,
    )

    numpy.testing.assert_equal(
        nx_tomo_sub_2.sample.rotation_angle[0:3],
        0.0,
    )
    numpy.testing.assert_array_equal(
        nx_tomo_sub_2.sample.rotation_angle[3:6],
        numpy.array([0.0, 10.5, 12.5]),
    )

    nx_tomo_sub_3 = NXtomo.sub_select_from_angle_offset(
        nx_tomo=nx_tomo,
        start_angle_offset=-10,
        angle_interval=300,
        shift_angles=False,
    )

    numpy.testing.assert_equal(
        nx_tomo_sub_3.instrument.detector.image_key_control,
        numpy.array(
            (
                ImageKey.DARK_FIELD,
                ImageKey.DARK_FIELD,
                ImageKey.FLAT_FIELD,
                ImageKey.PROJECTION,
                ImageKey.PROJECTION,
                ImageKey.PROJECTION,
                ImageKey.FLAT_FIELD,
                ImageKey.PROJECTION,
                ImageKey.PROJECTION,
                ImageKey.INVALID,
                ImageKey.FLAT_FIELD,
            )
        ),
    )

    nx_tomo_sub_4 = NXtomo.sub_select_from_angle_offset(
        nx_tomo=nx_tomo,
        start_angle_offset=-10,
        angle_interval=None,
        shift_angles=False,
    )
    numpy.testing.assert_equal(
        nx_tomo_sub_4.instrument.detector.image_key_control,
        numpy.array(
            (
                ImageKey.DARK_FIELD,
                ImageKey.DARK_FIELD,
                ImageKey.FLAT_FIELD,
                ImageKey.PROJECTION,
                ImageKey.PROJECTION,
                ImageKey.PROJECTION,
                ImageKey.FLAT_FIELD,
                ImageKey.PROJECTION,
                ImageKey.PROJECTION,
                ImageKey.PROJECTION,
                ImageKey.FLAT_FIELD,
            )
        ),
    )


def test_bliss_original_files(tmp_path):
    """
    test about NXtomo.bliss_original_files
    """
    test_dir = tmp_path / "test_bliss_original_files"
    test_dir.mkdir()

    nx_tomo_1 = NXtomo()
    with pytest.raises(TypeError):
        nx_tomo_1.bliss_original_files = 12

    nx_tomo_1.bliss_original_files = ("/path/1", "/path/2")

    nx_tomo_2 = NXtomo()
    nx_tomo_2.bliss_original_files = ("/path/2", "/path/3")

    nx_tomo_3 = NXtomo()

    nx_tomo_4 = NXtomo()
    nx_tomo_4.bliss_original_files = ()

    nx_tomo_concat = concatenate([nx_tomo_1, nx_tomo_2, nx_tomo_3])
    assert nx_tomo_concat.bliss_original_files == ("/path/1", "/path/2", "/path/3")

    output_nx_tomo_concat = os.path.join(test_dir, "nx_concat.nx")
    nx_tomo_concat.save(output_nx_tomo_concat, "/entry_concat")

    loaded_nx_tomo = NXtomo().load(output_nx_tomo_concat, "/entry_concat")
    assert loaded_nx_tomo.bliss_original_files == ("/path/1", "/path/2", "/path/3")

    output_nx_tomo_file = os.path.join(test_dir, "nx_tomo.nx")
    nx_tomo_3.save(output_nx_tomo_file, "/entry0000")
    loaded_nx_tomo = NXtomo().load(output_nx_tomo_file, "/entry0000")
    assert loaded_nx_tomo.bliss_original_files is None

    nx_tomo_4.save(output_nx_tomo_file, "/entry0000", overwrite=True)
    loaded_nx_tomo = NXtomo().load(output_nx_tomo_file, "/entry0000")
    assert loaded_nx_tomo.bliss_original_files == ()
