from unittest.mock import MagicMock, patch
import tempfile
import pytest
import os

import h5py

from dioptas.model.Configuration import Configuration
from dioptas.model.MapModel import MapModel
from dioptas.model.util.integration import iter_frames_sequential
from dioptas.tests.utility import unittest_data_path
import numpy as np

jcpds_path = os.path.join(unittest_data_path, "jcpds")
map_img_path = os.path.join(unittest_data_path, "map")
map_pattern_path = os.path.join(unittest_data_path, "map", "xy")
map_img_file_names = [
    f for f in os.listdir(map_img_path) if os.path.isfile(os.path.join(map_img_path, f))
]
map_img_file_paths = [
    os.path.join(map_img_path, filename) for filename in map_img_file_names
]

multi_file_img_path = os.path.join(
    unittest_data_path, "lambda", "testasapo1_1009_00002_m1_part00000.nxs"
)


@pytest.fixture
def configuration() -> Configuration:
    return Configuration()


@pytest.fixture
def map_model(configuration: Configuration) -> MapModel:
    return MapModel(configuration)


def test_create_map(map_model: MapModel, configuration: Configuration):
    configuration.calibration_model.load(
        os.path.join(unittest_data_path, "CeO2_Pilatus1M.poni")
    )
    map_model.load(map_img_file_paths)
    assert map_model.filepaths == map_img_file_paths
    assert len(map_model.pattern_intensities) == len(map_img_file_paths)
    assert map_model.dimension == (3, 3)
    assert map_model.map.shape == (3, 3)


def test_load_empty_filelist(map_model: MapModel, configuration: Configuration):
    with pytest.raises(ValueError):
        map_model.load([])


def test_load_files_with_different_dimensions(
    map_model: MapModel, configuration: Configuration
):
    file_paths = [
        os.path.join(unittest_data_path, "CeO2_Pilatus1M.tif"),
        os.path.join(unittest_data_path, "image_001.tif"),
    ]
    configuration.calibration_model.load(
        os.path.join(unittest_data_path, "CeO2_Pilatus1M.poni")
    )
    with pytest.raises(ValueError):
        map_model.load(file_paths)
    
    assert map_model.filepaths is None
    assert map_model.pattern_intensities is None
    assert map_model.dimension is None
    assert configuration.img_model.img_changed.blocked is False
    assert configuration.trim_trailing_zeros is True


def test_set_dimensions(map_model: MapModel, configuration: Configuration):
    configuration.calibration_model.load(
        os.path.join(unittest_data_path, "CeO2_Pilatus1M.poni")
    )
    map_model.load(map_img_file_paths[:6])
    assert len(map_model.pattern_intensities) == 6
    assert map_model.dimension == (2, 3)
    assert map_model.map.shape == (2, 3)

    map_model.set_dimension((3, 2))
    assert map_model.dimension == (3, 2)
    assert map_model.map.shape == (3, 2)

    map_model.set_dimension((1, 6))
    assert map_model.dimension == (1, 6)
    assert map_model.map.shape == (1, 6)


def test_set_wrong_dimensions(map_model: MapModel, configuration: Configuration):
    configuration.calibration_model.load(
        os.path.join(unittest_data_path, "CeO2_Pilatus1M.poni")
    )
    map_model.load(map_img_file_paths[:6])
    assert len(map_model.pattern_intensities) == 6
    assert map_model.dimension == (2, 3)
    assert map_model.map.shape == (2, 3)

    map_model.set_dimension((3, 3))
    assert map_model.dimension == (2, 3)
    assert map_model.map.shape == (2, 3)


def test_set_different_window(map_model: MapModel, configuration: Configuration):
    configuration.calibration_model.load(
        os.path.join(unittest_data_path, "CeO2_Pilatus1M.poni")
    )
    map_model.load(map_img_file_paths[:6])

    map_model.set_window((15, 16))
    assert map_model.window_intensities.all() > 0

    map_model.set_window((35, 40))  # window outside of pattern range
    assert map_model.window_intensities.all() == 0


def test_get_point_information(map_model: MapModel, configuration: Configuration):
    configuration.calibration_model.load(
        os.path.join(unittest_data_path, "CeO2_Pilatus1M.poni")
    )
    map_model.load(map_img_file_paths[:6])

    assert map_model.dimension == (2, 3)
    for i in range(6):
        column_index = i % 3
        row_index = i // 3
        point_info = map_model.get_point_info(row_index, column_index)
        assert point_info.filename == map_img_file_names[i]


def test_use_multi_file_img(map_model: MapModel, configuration: Configuration):
    configuration.calibration_model.load(
        os.path.join(unittest_data_path, "CeO2_Pilatus1M.poni")
    )
    map_model.load([multi_file_img_path])

    img_model = configuration.img_model

    assert map_model.filepaths == [multi_file_img_path]
    assert img_model.series_max == 10
    assert map_model.pattern_intensities.shape[0] == 10


def test_integrates_each_image_only_once(
    map_model: MapModel, configuration: Configuration
):
    configuration.calibration_model.load(
        os.path.join(unittest_data_path, "CeO2_Pilatus1M.poni")
    )
    x = np.linspace(0, 10, 100)
    y = np.sin(x)
    configuration.calibration_model.integrate_1d = MagicMock(return_value=(x, y))
    configuration.calibration_model.can_use_dioptrin_batch = MagicMock(return_value=False)
    map_model.load(map_img_file_paths)

    assert configuration.calibration_model.integrate_1d.call_count == len(
        map_img_file_paths
    )


def test_emits_point_integrated_signal(
    map_model: MapModel, configuration: Configuration
):
    configuration.calibration_model.load(
        os.path.join(unittest_data_path, "CeO2_Pilatus1M.poni")
    )
    # mock the integrate_1d method
    x = np.linspace(0, 10, 100)
    y = np.sin(x)
    configuration.calibration_model.integrate_1d = MagicMock(return_value=(x, y))

    listener = MagicMock()
    map_model.point_integrated.connect(listener)

    map_model.load(map_img_file_paths)
    assert listener.call_count == 9
    # assert listener.call_args_list == [(1), (2), (3), (4), (5), (6), (7), (8), (9)]


def test_emits_point_integrated_signal_with_multiimage_file(
    map_model: MapModel, configuration: Configuration
):
    configuration.calibration_model.load(
        os.path.join(unittest_data_path, "CeO2_Pilatus1M.poni")
    )
    # mock the integrate_1d method
    x = np.linspace(0, 10, 100)
    y = np.sin(x)
    configuration.calibration_model.integrate_1d = MagicMock(return_value=(x, y))

    listener = MagicMock()
    map_model.point_integrated.connect(listener)

    map_model.load([multi_file_img_path])
    assert listener.call_count == 10


def test_iter_frames_sequential_bitshuffle_shape_mismatch(configuration: Configuration):
    """Bitshuffle HDF5 path in iter_frames_sequential must reject shape mismatches."""
    img_model = configuration.img_model

    expected_shape = (2048, 2048)
    wrong_shape = (3262, 3108)

    mock_loader = MagicMock()
    mock_loader.gen_frames.return_value = iter([np.zeros(wrong_shape)])

    with patch(
        "dioptas.model.util.integration.try_open_bitshuffle_hdf5",
        return_value=mock_loader,
    ):
        gen = iter_frames_sequential(
            img_model,
            ["/fake/file.h5"],
            img_shape=expected_shape,
        )
        with pytest.raises(ValueError, match="expected"):
            next(gen)


def test_save_load_hdf5_round_trip(map_model: MapModel, configuration: Configuration):
    """Map state survives a save/load round-trip through HDF5."""
    configuration.calibration_model.load(
        os.path.join(unittest_data_path, "CeO2_Pilatus1M.poni")
    )
    map_model.load(map_img_file_paths)

    # tweak window and dimension so we verify they are restored
    map_model.set_window((15, 20))
    map_model.set_dimension((3, 3))

    with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as tmp:
        tmp_path = tmp.name

    try:
        # save
        with h5py.File(tmp_path, "w") as f:
            map_model.save_in_hdf5(f)

        # load into a fresh model
        new_model = MapModel(configuration)
        with h5py.File(tmp_path, "r") as f:
            new_model.load_from_hdf5(f)

        np.testing.assert_array_equal(new_model.pattern_x, map_model.pattern_x)
        np.testing.assert_array_equal(
            new_model.pattern_intensities, map_model.pattern_intensities
        )
        assert new_model.filepaths == map_model.filepaths
        assert len(new_model.point_infos) == len(map_model.point_infos)
        for a, b in zip(new_model.point_infos, map_model.point_infos):
            assert a.filepath == b.filepath
            assert a.frame_index == b.frame_index
        assert new_model.window == pytest.approx(map_model.window)
        assert new_model.dimension == map_model.dimension
        np.testing.assert_array_equal(new_model.map, map_model.map)
    finally:
        os.unlink(tmp_path)


def test_load_hdf5_without_map_group(map_model: MapModel):
    """Loading from an HDF5 file with no 'map' group is a no-op."""
    with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as tmp:
        tmp_path = tmp.name

    try:
        with h5py.File(tmp_path, "w") as f:
            pass  # empty file

        map_model.load_from_hdf5(h5py.File(tmp_path, "r"))
        assert map_model.filepaths is None
        assert map_model.map is None
    finally:
        os.unlink(tmp_path)
