from contextlib import nullcontext
from pathlib import Path
from typing import Literal

import numpy as np
import pytest
from numpy.typing import NDArray

from careamics.dataset_ng.image_stack.czi_image_stack import CziImageStack
from careamics.dataset_ng.image_stack.image_utils import channel_slice

# skip if fail imports
pylib = pytest.importorskip("pylibCZIrw")

from pylibCZIrw import czi as pyczi  # noqa: E402

T_EXPR = "does not contain a T axis"
Z_EXPR = "does not contain a Z axis"


def create_test_czi(file_path: Path, data: NDArray | list[NDArray]):
    if not isinstance(data, list):
        data = [data]

    with pyczi.create_czi(str(file_path)) as czi:
        xoffs = 0
        for scene_idx, scene_data in enumerate(data):
            for t in range(scene_data.shape[0]):
                for c in range(scene_data.shape[1]):
                    for z in range(scene_data.shape[2]):
                        czi.write(
                            scene_data[t, c, z],
                            plane={"C": c, "T": t, "Z": z},
                            location=(xoffs, 0),
                            scene=scene_idx,
                        )
            xoffs += scene_data.shape[-1] + 20


@pytest.mark.czi
class TestCziImageStack:

    @pytest.mark.parametrize(
        "orig_shape, depth_axis, expected_axes, expected_shape, sample_idx, "
        "expect_raise",
        [
            # 2-D Images
            ((1, 1, 1, 32, 48), "none", "SCYX", [1, 1, 32, 48], 0, nullcontext()),
            (
                (1, 1, 1, 32, 48),
                "T",
                "",
                [],
                0,
                pytest.raises(RuntimeError, match=T_EXPR),
            ),
            (
                (1, 1, 1, 32, 48),
                "Z",
                "",
                [],
                0,
                pytest.raises(RuntimeError, match=Z_EXPR),
            ),
            # 3-D Volumes
            ((1, 1, 16, 32, 48), "none", "SCYX", [16, 1, 32, 48], 9, nullcontext()),
            (
                (1, 1, 16, 32, 48),
                "T",
                "",
                [],
                9,
                pytest.raises(RuntimeError, match=T_EXPR),
            ),
            ((1, 1, 16, 32, 48), "Z", "SCZYX", [1, 1, 16, 32, 48], 0, nullcontext()),
            # 2-D Time-Series
            ((8, 1, 1, 32, 48), "none", "SCYX", [8, 1, 32, 48], 3, nullcontext()),
            (
                (8, 1, 1, 32, 48),
                "Z",
                "",
                [],
                3,
                pytest.raises(RuntimeError, match=Z_EXPR),
            ),
            ((8, 1, 1, 32, 48), "T", "SCTYX", [1, 1, 8, 32, 48], 0, nullcontext()),
            # 3-D Time-Series
            (
                (8, 1, 16, 32, 48),
                "none",
                "SCYX",
                [8 * 16, 1, 32, 48],
                35,
                nullcontext(),
            ),
            ((8, 1, 16, 32, 48), "Z", "SCZYX", [8, 1, 16, 32, 48], 7, nullcontext()),
            ((8, 1, 16, 32, 48), "T", "SCTYX", [16, 1, 8, 32, 48], 12, nullcontext()),
            # Multiple channels
            ((8, 3, 16, 32, 48), "Z", "SCZYX", [8, 3, 16, 32, 48], 2, nullcontext()),
        ],
    )
    def test_extract_patch(
        self,
        tmp_path: Path,
        orig_shape: tuple[int, ...],
        depth_axis: Literal["none", "Z", "T"],
        expected_axes: str,
        expected_shape: list[int],
        sample_idx: int,
        expect_raise,
    ):
        # reference data to compare against
        data = np.random.randn(*orig_shape).astype(np.float32)

        # save data as a czi file to ininitialise image stack with
        file_path = tmp_path / "test_czi.czi"
        create_test_czi(file_path=file_path, data=data)

        # initialise CziImageStack
        with expect_raise:
            image_stack = CziImageStack(data_path=file_path, depth_axis=depth_axis)

        # stop here if expecting an exception
        if expected_axes == "":
            return

        # check axes and shape
        assert image_stack.axes == expected_axes
        assert image_stack.data_shape == expected_shape

        # test extracted patch matches patch from reference data
        if len(expected_axes) < 5:
            coords = (11, 4)
            patch_size = (16, 9)

            extracted_patch = image_stack.extract_patch(
                sample_idx=sample_idx,
                channels=None,
                coords=coords,
                patch_size=patch_size,
            )

            data_ref = np.moveaxis(data, 2, 1)  # (T, Z, C, Y, X)
            data_ref = data_ref.reshape(-1, *data_ref.shape[-3:])
            patch_ref = data_ref[
                sample_idx,
                :,
                coords[0] : coords[0] + patch_size[0],
                coords[1] : coords[1] + patch_size[1],
            ]
        else:
            coords = (2, 11, 4)
            patch_size = (4, 16, 9)

            extracted_patch = image_stack.extract_patch(
                sample_idx=sample_idx,
                channels=None,
                coords=coords,
                patch_size=patch_size,
            )

            if "T" in expected_axes and expected_axes.index("T") == 2:
                data_ref = data.swapaxes(0, 2)
            else:
                data_ref = data
            patch_ref = data_ref[
                sample_idx,
                ...,
                coords[0] : coords[0] + patch_size[0],
                coords[1] : coords[1] + patch_size[1],
                coords[2] : coords[2] + patch_size[2],
            ]
        np.testing.assert_array_equal(extracted_patch, patch_ref)

    def test_multiple_scenes(
        self,
        tmp_path: Path,
    ):
        original_shapes = [
            [4, 1, 8, 16, 32],
            [4, 1, 8, 24, 18],
            [4, 1, 8, 45, 32],
        ]

        # reference data to compare against
        data_ref = [
            np.random.randn(*shape).astype(np.float32) for shape in original_shapes
        ]

        # save data as a czi file to ininitialise image stack with
        file_path = tmp_path / "test_czi.czi"
        create_test_czi(file_path=file_path, data=data_ref)

        # Test reading scene metadata
        scene_rectangles = CziImageStack.get_bounding_rectangles(file_path)
        assert len(scene_rectangles) == len(original_shapes)

        scene_rectangle_sizes = [(rect.h, rect.w) for rect in scene_rectangles.values()]
        expected_rectangle_sizes = [(shape[-2], shape[-1]) for shape in original_shapes]
        assert scene_rectangle_sizes == expected_rectangle_sizes

        # Test reading from individual scenes
        for scene_idx, expected_shape in enumerate(original_shapes):
            image_stack = CziImageStack(
                data_path=file_path, scene=scene_idx, depth_axis="Z"
            )
            assert image_stack.data_shape == expected_shape

            t = expected_shape[0] - scene_idx - 1
            coords = (2, 9, 4)
            patch_size = (4, 7, 13)

            extracted_patch = image_stack.extract_patch(
                sample_idx=t, channels=None, coords=coords, patch_size=patch_size
            )
            patch_ref = data_ref[scene_idx][
                t,
                :,
                coords[0] : coords[0] + patch_size[0],
                coords[1] : coords[1] + patch_size[1],
                coords[2] : coords[2] + patch_size[2],
            ]
            np.testing.assert_array_equal(extracted_patch, patch_ref)


@pytest.mark.czi
@pytest.mark.parametrize("axis", ["Z", "T"])
def test_z_padding(tmp_path: Path, axis):
    """Test that Z padding is applied to CziImageStack when requesting data outside
    the Z range."""
    if "Z" == axis:
        shape = (1, 2, 4, 4, 4)  # (T, C, Z, Y, X)
    else:
        shape = (4, 2, 1, 4, 4)  # (Z, C, T, Y, X)

    data = np.random.randn(*shape).astype(np.float32)

    file_path = tmp_path / "test_czi.czi"
    create_test_czi(file_path=file_path, data=data)

    # initialise CziImageStack
    image_stack = CziImageStack(data_path=file_path, depth_axis=axis)

    # extract patches
    sample_idx = 0
    channel_idx = 1
    patch_size = (4, 4, 4)
    z_length = shape[2] if "Z" == axis else shape[0]
    coordinates = [(-2, 0, 0), (2, 0, 0)]
    for coord in coordinates:
        patch = image_stack.extract_patch(
            sample_idx=sample_idx,
            channels=[channel_idx],
            coords=coord,
            patch_size=patch_size,
        )

        # build expected array
        z_start = max(coord[0], 0)
        z_end = min(coord[0] + patch_size[0], z_length)
        z_start_pad = z_start - coord[0]
        z_end_pad = z_end - coord[0]

        expected_patch = np.zeros(patch_size, dtype=np.float32)

        if "Z" == axis:
            expected_patch[z_start_pad:z_end_pad, :, :] = data[
                sample_idx, channel_idx, z_start:z_end, :, :
            ]
        else:  # T
            expected_patch[z_start_pad:z_end_pad, :, :] = data[
                z_start:z_end, channel_idx, sample_idx, :, :
            ]

        np.testing.assert_array_equal(patch[0], expected_patch)


@pytest.mark.czi
class TestCziImageStackChannels:

    @pytest.mark.parametrize(
        "shape, channels",
        [
            ((1, 1, 1, 32, 32), None),
            ((1, 1, 1, 32, 32), [0]),
            ((1, 3, 1, 32, 32), None),
            ((1, 3, 1, 32, 32), [0, 2]),
            ((1, 3, 1, 32, 32), [2]),
        ],
    )
    def test_extract_channels(
        self,
        tmp_path: Path,
        shape: tuple[int, ...],
        channels: list[int] | None,
    ):
        # reference data to compare against
        data = np.random.randn(*shape).astype(np.float32)

        # save data as a czi file to initialize image stack with
        file_path = tmp_path / "test_czi.czi"
        create_test_czi(file_path=file_path, data=data)

        # initialise CziImageStack
        image_stack = CziImageStack(data_path=file_path)

        # extract patch
        patch = image_stack.extract_patch(
            sample_idx=0,
            channels=channels,
            coords=(0, 0),
            patch_size=(8, 8),
        )
        assert len(patch.shape) == 3  # no Z
        assert (
            patch.shape[0] == len(channels) if channels is not None else data.shape[0]
        )

        expected_patch = data[
            0,
            channel_slice(channels),
            0,
            0 : 0 + 8,
            0 : 0 + 8,
        ]
        np.testing.assert_array_equal(patch, expected_patch)

    @pytest.mark.parametrize(
        "shape, channels",
        [
            ((2, 3, 1, 64, 64), [0, 4]),
            ((2, 3, 1, 64, 64), [3]),
        ],
    )
    def test_extract_channel_error(
        self,
        tmp_path: Path,
        shape: tuple[int, ...],
        channels: int,
    ):
        # reference data to compare against
        data = np.random.randn(*shape).astype(np.float32)

        # save data as a czi file to ininitialise image stack with
        file_path = tmp_path / "test_czi.czi"
        create_test_czi(file_path=file_path, data=data)

        # initialise CziImageStack
        image_stack = CziImageStack(data_path=file_path)

        expected_msg = (
            f"Channel index {channels[-1]} is out of bounds for data with "
            f"{shape[1]} channels. Check the provided `channels` "
            f"parameter in the configuration for erroneous channel "
            f"indices."
        )

        with pytest.raises(ValueError, match=expected_msg):
            image_stack.extract_patch(
                sample_idx=0,
                channels=channels,
                coords=(0, 0),
                patch_size=(16, 16),
            )
