# ------------------------------------------------------------------------
# RF-DETR
# Copyright (c) 2025 Roboflow. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
import pytest
import torch

from rfdetr.models.backbone.dinov2_with_windowed_attn import (
    Dinov2WithRegistersAttention,
    Dinov2WithRegistersSdpaAttention,
    WindowedDinov2WithRegistersBackbone,
    WindowedDinov2WithRegistersConfig,
    WindowedDinov2WithRegistersEmbeddings,
    WindowedDinov2WithRegistersModel,
    _find_pruneable_heads_and_indices,
    _get_aligned_output_features_output_indices,
)


def test_window_partition_forward_rectangular_preserves_shapes():
    """Regression test for WindowedDinov2WithRegistersEmbeddings.forward with rectangular input.

    Ensures window partitioning logic correctly handles H != W.
    """
    # Params: H_patches=6, W_patches=4, num_windows=2 -> 3x2 patches per window
    batch_size, hidden_size, patch_size, num_windows = 1, 64, 16, 2
    hp, wp, nr = 6, 4, 4
    h, w = hp * patch_size, wp * patch_size

    config = WindowedDinov2WithRegistersConfig(
        hidden_size=hidden_size,
        patch_size=patch_size,
        num_windows=num_windows,
        image_size=h,  # square image_size for positional embeddings
        num_register_tokens=nr,
    )
    model = WindowedDinov2WithRegistersEmbeddings(config)

    # Input is rectangular
    pixel_values = torch.randn(batch_size, 3, h, w)
    result = model(pixel_values)

    expected_batch = batch_size * (num_windows**2)
    expected_seq_len = 1 + nr + (hp // num_windows) * (wp // num_windows)

    assert result.shape == (expected_batch, expected_seq_len, hidden_size)


# Before fix in PR #448 the reshape used num_h_patches_per_window in both the height
# AND width dimension. This only fails when height and width produce different patch
# counts, so all tests below use non-square images (hp != wp).


@pytest.mark.parametrize(
    "hp, wp, num_windows",
    [
        (4, 6, 2),  # wider than tall
        (6, 4, 2),  # taller than wide
        (6, 9, 3),  # 3-window grid, non-square
        (8, 4, 2),  # 2:1 aspect ratio
    ],
)
def test_window_partition_nonsquare_does_not_raise(hp, wp, num_windows):
    """Before the fix, the reshape used num_h_patches_per_window for the width dimension, so the total element count
    mismatched and PyTorch raised a RuntimeError for any non-square image.

    The fix replaces that variable with num_w_patches_per_window, making the operation valid for all shapes.
    """
    hidden_size, patch_size, nr = 32, 16, 0
    h, w = hp * patch_size, wp * patch_size

    config = WindowedDinov2WithRegistersConfig(
        hidden_size=hidden_size,
        patch_size=patch_size,
        num_windows=num_windows,
        image_size=max(h, w),
        num_register_tokens=nr,
    )
    model = WindowedDinov2WithRegistersEmbeddings(config)
    pixel_values = torch.randn(1, 3, h, w)

    # This line would raise RuntimeError before the fix
    result = model(pixel_values)

    expected_batch = num_windows**2
    expected_seq_len = 1 + (hp // num_windows) * (wp // num_windows)
    assert result.shape == (expected_batch, expected_seq_len, hidden_size)


def test_window_partition_correct_window_content():
    """Verifies that after windowing each window contains the spatially correct patch tokens — not just that the shape
    is right.

    Layout with hp=4, wp=6, num_windows=2 (2x2 grid of windows):   Window (0,0): rows 0-1, cols 0-2   Window (0,1): rows
    0-1, cols 3-5   Window (1,0): rows 2-3, cols 0-2   Window (1,1): rows 2-3, cols 3-5

    Before the fix the reshape used num_h_patches_per_window for the width dim so it raised an error and never produced
    window content at all.
    """
    hidden_size, patch_size, num_windows, nr = 1, 16, 2, 0
    hp, wp = 4, 6
    h, w = hp * patch_size, wp * patch_size
    batch_size = 1

    config = WindowedDinov2WithRegistersConfig(
        hidden_size=hidden_size,
        patch_size=patch_size,
        num_windows=num_windows,
        image_size=max(h, w),
        num_register_tokens=nr,
        num_hidden_layers=1,
        num_attention_heads=1,
    )
    model = WindowedDinov2WithRegistersEmbeddings(config)

    # Disable position embeddings and cls token so we can track patch identity.
    # Each patch gets a unique value equal to its flat index (row * wp + col).
    with torch.no_grad():
        model.position_embeddings.zero_()
        model.cls_token.zero_()

    # Build a synthetic patch embedding: patch at (row, col) has value row*wp+col.
    # Shape after patch projection: (1, hp*wp, 1) — hidden_size=1 for simplicity.
    patch_ids = torch.arange(hp * wp, dtype=torch.float).view(1, hp * wp, 1)

    # Bypass the full forward pass and exercise the windowing logic directly.
    pixel_tokens = patch_ids  # (1, 24, 1)
    pixel_tokens_2d = pixel_tokens.view(batch_size, hp, wp, hidden_size)  # (1,4,6,1)

    num_h_patches_per_window = hp // num_windows  # 2
    num_w_patches_per_window = wp // num_windows  # 3

    # --- correct reshape (the fix) ---
    windowed = pixel_tokens_2d.reshape(
        batch_size * num_windows,
        num_h_patches_per_window,
        num_windows,
        num_w_patches_per_window,
        hidden_size,
    )
    windowed = windowed.permute(0, 2, 1, 3, 4)
    windowed = windowed.reshape(
        batch_size * num_windows**2,
        num_h_patches_per_window * num_w_patches_per_window,
        hidden_size,
    )

    # Expected content for each of the 4 windows (6 patches each):
    expected = torch.tensor(
        [
            # Window 0 (rows 0-1, cols 0-2): ids 0,1,2, 6,7,8
            [[0.0], [1.0], [2.0], [6.0], [7.0], [8.0]],
            # Window 1 (rows 0-1, cols 3-5): ids 3,4,5, 9,10,11
            [[3.0], [4.0], [5.0], [9.0], [10.0], [11.0]],
            # Window 2 (rows 2-3, cols 0-2): ids 12,13,14, 18,19,20
            [[12.0], [13.0], [14.0], [18.0], [19.0], [20.0]],
            # Window 3 (rows 2-3, cols 3-5): ids 15,16,17, 21,22,23
            [[15.0], [16.0], [17.0], [21.0], [22.0], [23.0]],
        ]
    )
    assert torch.equal(windowed, expected), f"Window content mismatch:\n{windowed}\n!=\n{expected}"


def test_buggy_reshape_raises_for_nonsquare():
    """Directly demonstrates what the pre-fix code did: using num_h_patches_per_window in the width position of the
    reshape causes a RuntimeError when the element count is not divisible by the (wrong) shape.

    With hidden_size=1 and hp=4, wp=6, num_windows=2 the total elements are 24 but the buggy target dims (2,2,2,2,-1)
    require a non-integer last dimension, so PyTorch raises RuntimeError.
    """
    hp, wp = 4, 6  # non-square: width > height
    num_windows = 2
    hidden_size = 1  # chosen so total / buggy-fixed-dims is non-integer

    num_h_patches_per_window = hp // num_windows  # 2
    num_w_patches_per_window = wp // num_windows  # 3
    batch_size = 1

    # Simulate pixel_tokens_with_pos_embed after the .view() call
    pixel_tokens_2d = torch.randn(batch_size, hp, wp, hidden_size)

    # The correct reshape (post-fix) must succeed
    pixel_tokens_2d.reshape(
        batch_size * num_windows,
        num_h_patches_per_window,
        num_windows,
        num_w_patches_per_window,  # correct
        hidden_size,
    )

    # The buggy reshape (pre-fix) must raise RuntimeError:
    # total elements = 1*4*6*1 = 24,  fixed-dims product = 2*2*2*2 = 16,  16 ∤ 24.
    with pytest.raises(RuntimeError):
        pixel_tokens_2d.reshape(
            batch_size * num_windows,
            num_h_patches_per_window,
            num_windows,
            num_h_patches_per_window,  # bug: height used for width
            -1,
        )


def test_buggy_reshape_silent_corruption_for_nonsquare():
    """When hidden_size happens to make the total element count divisible by the buggy target shape, PyTorch does NOT
    raise — instead the last dimension is inflated, which silently corrupts the tensor layout.

    Pre-fix with hp=4, wp=6, hidden_size=8, num_windows=2:   total elements = 1*4*6*8 = 192 buggy fixed dims = 2*2*2*2 =
    16  →  last dim inferred as 192/16 = 12 (not 8)

    The fix ensures the correct reshape always yields a last dim equal to hidden_size.
    """
    hp, wp = 4, 6
    num_windows = 2
    hidden_size = 8

    num_h_patches_per_window = hp // num_windows  # 2
    num_w_patches_per_window = wp // num_windows  # 3
    batch_size = 1

    pixel_tokens_2d = torch.randn(batch_size, hp, wp, hidden_size)

    # Buggy reshape silently infers last dim = 12 (not 8)
    buggy_out = pixel_tokens_2d.reshape(
        batch_size * num_windows,
        num_h_patches_per_window,
        num_windows,
        num_h_patches_per_window,  # bug
        -1,
    )
    assert buggy_out.shape[-1] != hidden_size, "Buggy reshape should produce wrong last dim"

    # Correct reshape always yields last dim == hidden_size
    correct_out = pixel_tokens_2d.reshape(
        batch_size * num_windows,
        num_h_patches_per_window,
        num_windows,
        num_w_patches_per_window,  # fix
        hidden_size,
    )
    assert correct_out.shape[-1] == hidden_size


# ---------------------------------------------------------------------------
# Tests for locally-copied utility functions (removed from transformers v5 public API)
# ---------------------------------------------------------------------------


class TestGetAlignedOutputFeaturesOutputIndices:
    """Tests for the local copy of get_aligned_output_features_output_indices."""

    def test_both_none_returns_last_stage(self):
        stage_names = ["stage1", "stage2", "stage3"]
        features, indices = _get_aligned_output_features_output_indices(None, None, stage_names)
        assert features == ["stage3"]
        assert indices == [2]

    def test_only_out_features_derives_indices(self):
        stage_names = ["stem", "layer1", "layer2", "layer3"]
        features, indices = _get_aligned_output_features_output_indices(["layer1", "layer3"], None, stage_names)
        assert features == ["layer1", "layer3"]
        assert indices == [1, 3]

    def test_only_out_indices_derives_features(self):
        stage_names = ["stem", "layer1", "layer2", "layer3"]
        features, indices = _get_aligned_output_features_output_indices(None, [0, 2], stage_names)
        assert features == ["stem", "layer2"]
        assert indices == [0, 2]

    def test_both_provided_returns_as_is(self):
        stage_names = ["stem", "layer1", "layer2"]
        features, indices = _get_aligned_output_features_output_indices(["layer1"], [1], stage_names)
        assert features == ["layer1"]
        assert indices == [1]

    def test_out_indices_converted_to_list(self):
        """out_indices supplied as a tuple must be returned as a list."""
        stage_names = ["stem", "layer1", "layer2"]
        _, indices = _get_aligned_output_features_output_indices(None, (1, 2), stage_names)
        assert isinstance(indices, list)
        assert indices == [1, 2]


class TestFindPruneableHeadsAndIndices:
    """Tests for the local copy of find_pruneable_heads_and_indices."""

    def test_no_pruning_returns_full_index(self):
        heads, index = _find_pruneable_heads_and_indices(set(), n_heads=4, head_size=3, already_pruned_heads=set())
        assert len(heads) == 0
        assert len(index) == 12  # 4 * 3, nothing masked

    @pytest.mark.parametrize(
        "head_to_prune, expected_index",
        [
            pytest.param({0}, list(range(3, 12)), id="prune-first-head"),
            pytest.param({3}, list(range(9)), id="prune-last-head"),
        ],
    )
    def test_prune_single_head_removes_correct_rows(self, head_to_prune, expected_index):
        # Head N masked → N*head_size indices removed; remaining = n_heads*head_size - head_size = 9
        heads, index = _find_pruneable_heads_and_indices(
            head_to_prune, n_heads=4, head_size=3, already_pruned_heads=set()
        )
        assert heads == head_to_prune
        assert len(index) == 9
        assert index.tolist() == expected_index

    def test_already_pruned_head_adjusts_offset(self):
        # Head 0 was already pruned. Now pruning head 1 (which is now effective head 0
        # after offset adjustment) should remove 3 more indices from the effective mask.
        heads, index = _find_pruneable_heads_and_indices({1}, n_heads=4, head_size=3, already_pruned_heads={0})
        assert 1 in heads
        assert len(index) == 9  # 4*3 - 3 pruned


# ---------------------------------------------------------------------------
# Smoke tests for WindowedDinov2WithRegistersBackbone
# ---------------------------------------------------------------------------


def _minimal_backbone_config(**kwargs) -> WindowedDinov2WithRegistersConfig:
    """Return the smallest valid config for backbone instantiation tests.

    Examples:
        >>> cfg = _minimal_backbone_config(hidden_size=32)
        >>> cfg.hidden_size
        32
    """
    defaults = dict(
        hidden_size=32,
        num_hidden_layers=1,
        num_attention_heads=2,
        intermediate_size=64,
        patch_size=16,
        image_size=64,
        num_register_tokens=0,
        num_windows=1,
    )
    defaults.update(kwargs)
    return WindowedDinov2WithRegistersConfig(**defaults)


class TestWindowedDinov2WithRegistersBackbone:
    """Smoke tests that guard against _init_transformers_backbone() API regressions."""

    @pytest.mark.parametrize(
        "attr",
        [
            pytest.param("stage_names", id="stage_names"),
            pytest.param("out_features", id="out_features"),
        ],
    )
    def test_instantiation_sets_list_attribute(self, attr):
        config = _minimal_backbone_config()
        backbone = WindowedDinov2WithRegistersBackbone(config)
        assert hasattr(backbone, attr)
        assert isinstance(getattr(backbone, attr), list)
        assert len(getattr(backbone, attr)) > 0

    def test_forward_returns_backbone_output(self):
        config = _minimal_backbone_config()
        backbone = WindowedDinov2WithRegistersBackbone(config)
        backbone.eval()
        pixel_values = torch.randn(1, 3, 64, 64)
        with torch.no_grad():
            output = backbone(pixel_values)
        assert hasattr(output, "feature_maps")
        assert len(output.feature_maps) == len(backbone.out_features)


# ---------------------------------------------------------------------------
# Test for output_attentions=True SDPA fallback path
# ---------------------------------------------------------------------------


class TestSdpaFallbackWithOutputAttentions:
    """Guards the output_attentions behaviour in windowed attention."""

    def test_output_attentions_true_raises(self):
        """Windowed attention explicitly does not support output_attentions=True."""
        config = _minimal_backbone_config()
        model = WindowedDinov2WithRegistersModel(config)
        model.eval()
        pixel_values = torch.randn(1, 3, 64, 64)
        with torch.no_grad():
            with pytest.raises(AssertionError, match="output_attentions is not supported for windowed attention"):
                model(pixel_values, output_attentions=True)


class TestSetAttnImplementation:
    """Tests for WindowedDinov2WithRegistersModel.set_attn_implementation."""

    @pytest.mark.parametrize(
        "switches, expected_impl, expected_cls",
        [
            pytest.param(["eager"], "eager", Dinov2WithRegistersAttention, id="sdpa-to-eager"),
            pytest.param(["eager", "sdpa"], "sdpa", Dinov2WithRegistersSdpaAttention, id="roundtrip-back-to-sdpa"),
        ],
    )
    def test_switch_updates_config_and_layers(self, switches, expected_impl, expected_cls):
        """After each call in *switches*, config and all layer attention modules reflect the final impl."""
        config = _minimal_backbone_config()
        model = WindowedDinov2WithRegistersModel(config)

        for impl in switches:
            model.set_attn_implementation(impl)

        assert model.config._attn_implementation == expected_impl
        for layer in model.encoder.layer:
            assert type(layer.attention) is expected_cls

    def test_invalid_implementation_raises(self):
        """Passing an unknown key raises ValueError with a clear message."""
        config = _minimal_backbone_config()
        model = WindowedDinov2WithRegistersModel(config)

        with pytest.raises(ValueError, match="Unknown attn_implementation"):
            model.set_attn_implementation("flash_attention_2")


@pytest.mark.parametrize(
    "h, w, num_windows, should_raise",
    [
        pytest.param(64, 64, 2, False, id="valid-square"),
        pytest.param(64, 96, 2, False, id="valid-rectangular"),
        pytest.param(32, 32, 1, False, id="num_windows-1-valid"),
        pytest.param(33, 64, 2, True, id="h-not-divisible"),
        pytest.param(64, 33, 2, True, id="w-not-divisible"),
        pytest.param(33, 33, 2, True, id="both-not-divisible"),
    ],
)
def test_forward_validates_spatial_dims(h: int, w: int, num_windows: int, should_raise: bool) -> None:
    """WindowedDinov2WithRegistersEmbeddings raises ValueError for incompatible dims.

    Both H and W must be divisible by patch_size * num_windows.  The check must survive Python's -O flag (assert would
    be silently stripped).
    """
    patch_size = 16
    config = WindowedDinov2WithRegistersConfig(
        hidden_size=32,
        patch_size=patch_size,
        num_windows=num_windows,
        image_size=max(h, w),
        num_register_tokens=0,
    )
    model = WindowedDinov2WithRegistersEmbeddings(config)
    pixel_values = torch.randn(1, 3, h, w)
    if should_raise:
        with pytest.raises(ValueError, match="divisible"):
            model(pixel_values)
    else:
        model(pixel_values)  # must not raise


def _make_small_model() -> WindowedDinov2WithRegistersModel:
    """Return the smallest valid WindowedDinov2WithRegistersModel for unit tests.

    Examples:
        >>> model = _make_small_model()
        >>> model.config.hidden_size
        32
    """
    config = WindowedDinov2WithRegistersConfig(
        hidden_size=32,
        num_hidden_layers=2,
        num_attention_heads=4,
        image_size=32,
        patch_size=16,
        num_register_tokens=2,
    )
    return WindowedDinov2WithRegistersModel(config)


class TestSetAttnImplementationPreservesWeights:
    """set_attn_implementation must transfer trained weights to the new attention module.

    Before the fix the method replaced each layer's attention with a freshly constructed (randomly initialised) module,
    silently discarding all trained q/k/v/output weights.
    """

    @pytest.mark.parametrize(
        "from_impl, to_impl",
        [
            pytest.param("sdpa", "eager", id="sdpa_to_eager"),
            pytest.param("eager", "sdpa", id="eager_to_sdpa"),
        ],
    )
    def test_query_weight_preserved_after_switch(self, from_impl: str, to_impl: str) -> None:
        """After switching implementation the query weight tensor must be unchanged."""
        model = _make_small_model()
        model.set_attn_implementation(from_impl)

        # Record the query weights of every layer before switching.
        before = [layer.attention.attention.query.weight.clone() for layer in model.encoder.layer]

        model.set_attn_implementation(to_impl)

        after = [layer.attention.attention.query.weight for layer in model.encoder.layer]
        for layer_idx, (w_before, w_after) in enumerate(zip(before, after)):
            assert torch.equal(w_before, w_after), (
                f"Layer {layer_idx}: query weight changed after set_attn_implementation({from_impl!r} → {to_impl!r})"
            )

    def test_key_and_value_weights_preserved(self) -> None:
        """Key and value weights must also survive the implementation switch."""
        model = _make_small_model()
        model.set_attn_implementation("sdpa")

        key_before = [layer.attention.attention.key.weight.clone() for layer in model.encoder.layer]
        val_before = [layer.attention.attention.value.weight.clone() for layer in model.encoder.layer]

        model.set_attn_implementation("eager")

        for layer_idx, layer in enumerate(model.encoder.layer):
            assert torch.equal(key_before[layer_idx], layer.attention.attention.key.weight), (
                f"Layer {layer_idx}: key weight changed after implementation switch"
            )
            assert torch.equal(val_before[layer_idx], layer.attention.attention.value.weight), (
                f"Layer {layer_idx}: value weight changed after implementation switch"
            )

    def test_output_dense_weight_preserved(self) -> None:
        """The output projection (dense) weight must survive the implementation switch."""
        model = _make_small_model()
        model.set_attn_implementation("sdpa")

        dense_before = [layer.attention.output.dense.weight.clone() for layer in model.encoder.layer]

        model.set_attn_implementation("eager")

        for layer_idx, layer in enumerate(model.encoder.layer):
            assert torch.equal(dense_before[layer_idx], layer.attention.output.dense.weight), (
                f"Layer {layer_idx}: output dense weight changed after implementation switch"
            )

    def test_config_updated_after_switch(self) -> None:
        """config._attn_implementation must reflect the new implementation after switching."""
        model = _make_small_model()
        model.set_attn_implementation("eager")
        assert model.config._attn_implementation == "eager"
        model.set_attn_implementation("sdpa")
        assert model.config._attn_implementation == "sdpa"

    def test_attention_module_type_after_switch(self) -> None:
        """After switching to eager, every layer must hold a non-SDPA attention class."""
        model = _make_small_model()
        model.set_attn_implementation("eager")
        for layer in model.encoder.layer:
            assert isinstance(layer.attention, Dinov2WithRegistersAttention)
            assert not isinstance(layer.attention, Dinov2WithRegistersSdpaAttention)

    def test_invalid_implementation_raises_value_error(self) -> None:
        """An unknown implementation name must raise ValueError before touching any layer."""
        model = _make_small_model()
        with pytest.raises(ValueError, match="Unknown attn_implementation"):
            model.set_attn_implementation("flash_attn")
