import numpy as np
import pytest

from napari.components import Dims
from napari.layers import Image, Labels
from napari.layers._scalar_field.scalar_field import ScalarFieldBase
from napari.utils._test_utils import (
    validate_all_params_in_docstring,
    validate_docstring_parent_class_consistency,
    validate_kwargs_sorted,
)


def test_docstring():
    validate_all_params_in_docstring(ScalarFieldBase)
    validate_kwargs_sorted(ScalarFieldBase)
    validate_docstring_parent_class_consistency(
        ScalarFieldBase, skip=('data', 'ndim', 'multiscale')
    )


def test_multiscale_thumbnail_level_prematerialized():
    """Thumbnail level materializes as numpy via the level materializer."""
    rng = np.random.default_rng(0)
    data = [rng.random((128, 128)), rng.random((64, 64))]
    layer = Image(data, multiscale=True)
    result = layer._level_materializer(layer._thumbnail_level)
    assert isinstance(result, np.ndarray)
    np.testing.assert_array_equal(result, data[layer._thumbnail_level])


def test_3d_multiscale_thumbnail_not_prematerialized():
    """3D multiscale images should not create a level materializer: volumes can be too large."""
    data = [np.random.random((16, 128, 128)), np.random.random((8, 64, 64))]
    layer = Image(data, multiscale=True)
    assert layer._level_materializer is None


def test_single_scale_no_prematerialization():
    """Single-scale images should not materialize and cache."""
    layer = Image(np.random.random((32, 32)))
    assert layer._level_materializer is None


def test_multiscale_slice_request_receives_preselected_sources():
    """Slice requests should receive the preselected sources."""
    data = [np.random.random((64, 64)), np.random.random((32, 32))]
    layer = Image(data, multiscale=True)
    dims = Dims(
        ndim=2,
        ndisplay=2,
        range=tuple((0, s - 1, 1) for s in data[0].shape),
    )
    request = layer._slicing_state._make_slice_request(dims)
    assert layer.data_level == layer._thumbnail_level
    # slice request should re-use thumbnail source when possible
    assert request.data_at_data_level is request.data_at_thumbnail_level
    np.testing.assert_array_equal(
        request.data_at_thumbnail_level, data[layer._thumbnail_level]
    )


def test_multiscale_slice_request_keeps_non_thumbnail_image_source():
    """Non-thumbnail views should keep their own level."""
    data = [np.random.random((64, 64)), np.random.random((32, 32))]
    layer = Image(data, multiscale=True)
    layer.data_level = 0
    dims = Dims(
        ndim=2,
        ndisplay=2,
        range=tuple((0, s - 1, 1) for s in data[0].shape),
    )
    request = layer._slicing_state._make_slice_request(dims)
    assert request.data_at_data_level is data[0]
    np.testing.assert_array_equal(
        request.data_at_thumbnail_level, data[layer._thumbnail_level]
    )


def test_thumbnail_level_data_refreshed_on_data_replacement():
    """Replacing layer.data must bind a fresh materializer for the new data
    so slices never read from the old dataset."""
    data1 = [np.zeros((64, 64)), np.zeros((32, 32))]
    data2 = [np.ones((90, 90)), np.ones((45, 45)), np.ones((22, 22))]

    layer = Image(data1, multiscale=True)
    assert layer._thumbnail_level == 1
    np.testing.assert_array_equal(
        layer._level_materializer(layer._thumbnail_level), data1[1]
    )

    layer.data = data2
    assert layer._thumbnail_level == 2  # new last level
    np.testing.assert_array_equal(
        layer._level_materializer(layer._thumbnail_level), data2[2]
    )


def test_thumbnail_level_data_uses_new_data_ndim_2d_to_3d():
    """Replacing 2D data with 3D should set materializer to None — 3D volumes
    are left as lazy data-array references."""
    from napari.layers._multiscale_data import MultiScaleData

    data2d = [np.zeros((64, 64)), np.zeros((32, 32))]
    data3d = [np.ones((16, 64, 64)), np.ones((8, 32, 32))]

    layer = Image(data2d, multiscale=True)
    assert layer._level_materializer is not None  # 2D: materializer present

    layer._data = MultiScaleData(data3d)
    layer._reset_thumbnail_level_data()

    assert layer._level_materializer is None  # 3D: no materializer


def test_thumbnail_level_data_uses_new_data_ndim_3d_to_2d():
    """Replacing 3D data with 2D should produce a fresh materializer for
    the new 2D data."""
    from napari.layers._multiscale_data import MultiScaleData

    data3d = [np.zeros((16, 64, 64)), np.zeros((8, 32, 32))]
    data2d = [np.ones((64, 64)), np.ones((32, 32))]

    layer = Image(data3d, multiscale=True)
    old_materializer = layer._level_materializer

    layer._data = MultiScaleData(data2d)
    layer._reset_thumbnail_level_data()
    assert layer._level_materializer is not old_materializer
    np.testing.assert_array_equal(
        layer._level_materializer(layer._thumbnail_level),
        data2d[layer._thumbnail_level],
    )


def test_rgb_2d_multiscale_prematerialized():
    """RGB (H, W, 3) multiscale images materialise correctly via the closure."""
    data = [np.ones((128, 128, 3)), np.ones((64, 64, 3))]
    layer = Image(data, multiscale=True, rgb=True)

    assert layer.ndim == 2  # spatial dims only
    tld = layer._level_materializer(layer._thumbnail_level)
    assert isinstance(tld, np.ndarray)
    assert tld.shape == (64, 64, 3)
    np.testing.assert_array_equal(tld, data[1])


@pytest.mark.parametrize('Layer', [Image, Labels])
def test_data_setter_updates_transforms(Layer):
    """Replacing data with different ndim should expand transforms."""
    data_2d = np.zeros((10, 10), dtype=np.uint8)
    layer = Layer(data_2d, scale=(2, 3))
    assert layer.ndim == 2
    assert len(layer.scale) == 2

    data_3d = np.zeros((5, 10, 10), dtype=np.uint8)
    layer.data = data_3d
    assert layer.ndim == 3
    assert len(layer.scale) == 3


@pytest.mark.parametrize('Layer', [Image, Labels])
def test_invisible_layer_keeps_placeholder_rank_in_3d(Layer):
    """When ndisplay flips while the layer is invisible, the cached
    placeholder must still match ndisplay's rank, otherwise vispy crashes
    with "Volume visual needs a 3D array." when it later reads
    layer._data_view.
    """
    layer = Layer(np.zeros((4, 5, 6), dtype=np.uint8), visible=False)
    layer._slice_dims(Dims(ndim=3, ndisplay=3))
    assert layer._data_view.ndim == 3
    layer.visible = True
    layer.visible = False
    assert layer._data_view.ndim == 3


# ---------------------------------------------------------------------------
# locked_data_level tests
# ---------------------------------------------------------------------------

_MULTISCALE_SHAPES_2D = [(4000, 3000), (2000, 1500), (1000, 750), (500, 375)]
_MULTISCALE_SHAPES_3D = [(8, 40, 20), (4, 20, 10), (2, 10, 5)]


def _make_multiscale_2d():
    """2D multiscale data where each level has a distinct constant value."""
    return [
        np.full(s, fill_value=(i + 1) * 10, dtype=np.uint8)
        for i, s in enumerate(_MULTISCALE_SHAPES_2D)
    ]


def _make_multiscale_3d():
    """3D multiscale data where each level has a distinct constant value."""
    return [
        np.full(s, fill_value=(i + 1) * 10, dtype=np.uint8)
        for i, s in enumerate(_MULTISCALE_SHAPES_3D)
    ]


@pytest.mark.parametrize('Layer', [Image, Labels])
class TestLockedDataLevel:
    """Tests for the locked_data_level property on ScalarFieldBase."""

    def test_locked_data_level_property(self, Layer):
        """Default is None; setting valid/invalid values works correctly."""
        data = _make_multiscale_3d()
        layer = Layer(data, multiscale=True)

        assert layer.locked_data_level is None

        layer.locked_data_level = 1
        assert layer.locked_data_level == 1

        layer.locked_data_level = None
        assert layer.locked_data_level is None

        # Out-of-range values are silently ignored (e.g. for linked layers
        # where one layer may have fewer levels than another).
        layer.locked_data_level = 1
        layer.locked_data_level = len(data)
        assert layer.locked_data_level == 1  # unchanged

        layer.locked_data_level = -1
        assert layer.locked_data_level == 1  # unchanged

    def test_locked_data_level_reset_on_data_change(self, Layer):
        """Setting .data resets locked_data_level to None (auto)."""
        data = _make_multiscale_3d()
        layer = Layer(data, multiscale=True)

        # Lock to the last level
        layer.locked_data_level = len(data) - 1
        assert layer.locked_data_level == len(data) - 1

        # Replace with data that has fewer levels
        fewer_levels = data[:2]
        layer.data = fewer_levels

        # The old locked level would be out of range; it must be reset
        assert layer.locked_data_level is None


def _draw_layer(layer, shape_threshold=(800, 600)):
    """Call ``_update_draw`` with a viewport that sees the full data extent."""
    displayed = layer._slice_input.displayed
    corner_pixels = layer._extent_data[:, displayed]
    layer._update_draw(
        scale_factor=1,
        corner_pixels_displayed=corner_pixels,
        shape_threshold=shape_threshold,
    )


@pytest.mark.parametrize('Layer', [Image, Labels])
class TestLockedDataLevelDraw:
    """Tests for locked_data_level interacting with the draw cycle.

    These call ``_update_draw`` and ``_slice_dims`` directly instead of
    going through a full Qt viewer, so they run in headless CI
    environments without Qt bindings.
    """

    def test_lock_overrides_auto_2d(self, Layer):
        """Locking to level 0 keeps full resolution even when the
        viewport is wide enough that auto would pick a coarser level."""
        data = _make_multiscale_2d()
        layer = Layer(data, multiscale=True)

        _draw_layer(layer)
        auto_level = layer.data_level
        assert auto_level > 0, 'sanity: auto should pick a coarser level'

        layer.locked_data_level = 0
        _draw_layer(layer)

        assert layer.data_level == 0

    def test_lock_overrides_auto_3d(self, Layer):
        """In 3D, auto defaults to the coarsest level; locking to 0
        should override that."""
        data = _make_multiscale_3d()
        layer = Layer(data, multiscale=True)

        layer._slice_dims(Dims(ndim=3, ndisplay=3))
        _draw_layer(layer)
        assert layer.data_level == len(data) - 1, (
            'sanity: 3D auto should pick coarsest'
        )

        layer.locked_data_level = 0
        _draw_layer(layer)

        assert layer.data_level == 0

    def test_unlock_restores_auto_2d(self, Layer):
        """Setting locked_data_level back to None restores automatic
        level selection in 2D."""
        data = _make_multiscale_2d()
        layer = Layer(data, multiscale=True)

        _draw_layer(layer)
        auto_level = layer.data_level

        layer.locked_data_level = 0
        _draw_layer(layer)
        assert layer.data_level == 0

        layer.locked_data_level = None
        _draw_layer(layer)
        assert layer.data_level == auto_level

    def test_unlock_restores_auto_3d(self, Layer):
        """Setting locked_data_level back to None restores coarsest-level
        default in 3D."""
        data = _make_multiscale_3d()
        layer = Layer(data, multiscale=True)

        layer._slice_dims(Dims(ndim=3, ndisplay=3))
        layer.locked_data_level = 0
        _draw_layer(layer)
        assert layer.data_level == 0

        layer.locked_data_level = None
        _draw_layer(layer)
        assert layer.data_level == len(data) - 1

    def test_ndisplay_transitions(self, Layer):
        """Lock and auto behaviour across 2D/3D switches."""
        data = _make_multiscale_3d()
        layer = Layer(data, multiscale=True)
        coarsest = len(data) - 1

        # -- auto 2D->3D: should pick coarsest
        _draw_layer(layer)
        layer._slice_dims(Dims(ndim=3, ndisplay=3))
        _draw_layer(layer)
        assert layer.data_level == coarsest

        # -- auto 3D->2D: should leave coarsest-level default behind
        layer._slice_dims(Dims(ndim=3, ndisplay=2))
        _draw_layer(layer)
        assert layer.data_level != coarsest

        # -- locked 2D->3D: lock should persist
        layer.locked_data_level = 0
        _draw_layer(layer)
        assert layer.data_level == 0

        layer._slice_dims(Dims(ndim=3, ndisplay=3))
        _draw_layer(layer)
        assert layer.data_level == 0

        # -- locked 3D->2D: lock should persist
        layer._slice_dims(Dims(ndim=3, ndisplay=2))
        _draw_layer(layer)
        assert layer.data_level == 0

    def test_locked_level_refreshes_on_stale_data_level(self, Layer):
        """_update_draw refreshes when _data_level drifts from locked."""
        data = _make_multiscale_3d()
        layer = Layer(data, multiscale=True)

        layer._slice_dims(Dims(ndim=3, ndisplay=3))
        layer.locked_data_level = 0
        _draw_layer(layer)
        assert layer.data_level == 0

        # Simulate _data_level drifting out of sync (e.g. internal reset)
        layer._data_level = len(data) - 1
        _draw_layer(layer)
        # _update_draw should have corrected it back to the locked level
        assert layer.data_level == 0

    def test_locked_level_loads_correct_data_3d(self, Layer):
        """Verify the actual pixel values match the locked level."""
        data = _make_multiscale_3d()
        layer = Layer(data, multiscale=True)

        layer._slice_dims(Dims(ndim=3, ndisplay=3))
        _draw_layer(layer)

        for level_idx in range(len(data)):
            expected_value = (level_idx + 1) * 10
            layer.locked_data_level = level_idx
            layer._slice_dims(Dims(ndim=3, ndisplay=3))
            view = layer._data_view
            assert np.all(view == expected_value), (
                f'Level {level_idx}: expected all {expected_value}, '
                f'got unique values {np.unique(view)}'
            )


def test_set_view_slice_3d_multiscale_corners_without_viewer():
    """Switching a multiscale layer to ndisplay=3 without a viewer
    (no _update_draw) must set corner_pixels and _data_level to the
    coarsest level so slicing produces correct 3D data.

    Regression test: without the fix, corner_pixels stayed at the 2D
    extent and _data_level was not updated, leading to wrong or empty
    slices.
    """
    data = [
        np.full((8, 40, 20), 10, dtype=np.uint8),
        np.full((4, 20, 10), 20, dtype=np.uint8),
        np.full((2, 10, 5), 30, dtype=np.uint8),
    ]
    layer = Image(data, multiscale=True)

    dims_3d = Dims(ndim=3, ndisplay=3)
    layer._slice_dims(dims_3d)

    coarsest = len(data) - 1
    assert layer.data_level == coarsest, (
        f'Expected data_level={coarsest}, got {layer.data_level}'
    )

    expected_shape = np.array(data[coarsest].shape)
    corners = layer.corner_pixels
    actual_extent = corners[1, :] - corners[0, :] + 1
    np.testing.assert_array_equal(
        actual_extent,
        expected_shape,
        err_msg=(
            'corner_pixels should span the full coarsest level '
            'when switching to 3D without a viewer'
        ),
    )
