# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Shape broadcasting regression tests for RotaryEmbedding.forward_native.

Covers the batch-dimension shape mismatch fix: forward_native now strips the
batch dim from 3D cos/sin before applying RoPE, consistent with forward_cuda,
forward_hip, and apply_rotary_emb_mindiesd paths.
"""

from __future__ import annotations

import sys
import types

import pytest
import torch

from vllm_omni.diffusion.layers.rope import (
    RotaryEmbedding,
    RotaryEmbeddingWan,
    apply_rotary_emb_mindiesd,
)

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


def _make_inputs(
    batch: int,
    seq_len: int,
    num_heads: int,
    head_dim: int,
    cos_batch: int | None = None,
    dtype: torch.dtype = torch.float32,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Generate x, cos, sin tensors for RoPE testing.

    Args:
        batch: Batch dimension of x.
        seq_len: Sequence length.
        num_heads: Number of attention heads.
        head_dim: Head dimension (must be even).
        cos_batch: Batch dimension of cos/sin. Defaults to batch if None.
        dtype: Data type for the tensors.
    """
    if cos_batch is None:
        cos_batch = batch

    torch.manual_seed(42)
    x = torch.randn(batch, seq_len, num_heads, head_dim, dtype=dtype)
    cos = torch.randn(cos_batch, seq_len, head_dim // 2, dtype=dtype)
    sin = torch.randn(cos_batch, seq_len, head_dim // 2, dtype=dtype)
    return x, cos, sin


class TestRotaryEmbeddingNativeShapeRegression:
    """Regression tests for forward_native shape handling (#8e297d81)."""

    @pytest.fixture
    def rope_neox(self) -> RotaryEmbedding:
        """NeoX-style (non-interleaved) RoPE."""
        return RotaryEmbedding(is_neox_style=True)

    @pytest.fixture
    def rope_gptj(self) -> RotaryEmbedding:
        """GPT-J style (interleaved) RoPE."""
        return RotaryEmbedding(is_neox_style=False)

    # ── 2D cos/sin (no batch dim) ──────────────────────────────────────

    def test_2d_cos_sin_unchanged(self, rope_neox: RotaryEmbedding) -> None:
        """2D cos/sin [S, D/2] pass through unchanged (no-op path)."""
        x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, cos_batch=2)
        # Manually drop batch dim to create 2D cos/sin
        cos_2d = cos[0]
        sin_2d = sin[0]
        assert cos_2d.dim() == 2
        output = rope_neox.forward_native(x, cos_2d, sin_2d)
        assert output.shape == x.shape
        assert not torch.isnan(output).any()
        assert not torch.isinf(output).any()

    # ── 3D cos/sin — batch-dim stripping ──────────────────────────────

    @pytest.mark.parametrize("cos_batch", [1, 2, 4])
    def test_3d_cos_sin_no_error(self, rope_neox: RotaryEmbedding, cos_batch: int) -> None:
        """3D cos/sin with batch dim should not raise errors."""
        x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, cos_batch=cos_batch)
        output = rope_neox.forward_native(x, cos, sin)
        assert output.shape == x.shape

    @pytest.mark.parametrize("interleaved", [False, True])
    def test_3d_cos_identical_to_2d(self, interleaved: bool) -> None:
        """3D cos/sin with identical batch slices == 2D cos/sin result."""
        rope = RotaryEmbedding(is_neox_style=not interleaved)
        B, S, H, D = 2, 16, 8, 64
        x = torch.randn(B, S, H, D)
        cos_2d = torch.randn(S, D // 2)
        sin_2d = torch.randn(S, D // 2)

        # 3D variant: expand the same 2D data across batch
        cos_3d = cos_2d.unsqueeze(0).expand(B, -1, -1)
        sin_3d = sin_2d.unsqueeze(0).expand(B, -1, -1)

        out_2d = rope.forward_native(x, cos_2d, sin_2d)
        out_3d = rope.forward_native(x, cos_3d, sin_3d)

        torch.testing.assert_close(out_3d, out_2d, atol=0, rtol=0)

    # ── Batch size mismatch — the regression scenario ─────────────────

    @pytest.mark.parametrize(
        "x_batch,cos_batch",
        [
            (2, 1),  # cos has fewer batch elements
            (1, 2),  # cos has more batch elements (takes [0])
            (4, 1),  # large x batch, single cos
        ],
    )
    def test_batch_size_mismatch_regression(
        self,
        rope_neox: RotaryEmbedding,
        x_batch: int,
        cos_batch: int,
    ) -> None:
        """forward_native works when cos batch dim differs from x batch dim.

        This is the exact regression scenario — before the fix, a 3D cos/sin
        with a different batch size than x would cause a shape mismatch error
        inside apply_rotary_emb_torch during broadcast.
        """
        x, cos, sin = _make_inputs(batch=x_batch, seq_len=16, num_heads=8, head_dim=64, cos_batch=cos_batch)
        output = rope_neox.forward_native(x, cos, sin)
        assert output.shape == x.shape
        assert output.dtype == x.dtype

    # ── Batch consistency — all elements get same rotation ────────────

    def test_2d_cos_broadcasts_identically_across_batch(self, rope_neox: RotaryEmbedding) -> None:
        """With 2D cos, identical x slices → identical output slices."""
        B, S, H, D = 4, 16, 8, 64
        x_template = torch.randn(1, S, H, D)
        x = x_template.expand(B, -1, -1, -1)  # all batch elements identical
        cos = torch.randn(S, D // 2)
        sin = torch.randn(S, D // 2)

        output = rope_neox.forward_native(x, cos, sin)

        # All output batch slices must be identical
        for i in range(1, B):
            torch.testing.assert_close(output[i], output[0], atol=0, rtol=0)

    def test_3d_identical_cos_broadcasts_identically_across_batch(self, rope_neox: RotaryEmbedding) -> None:
        """With 3D cos (identical slices), identical x → identical output."""
        B, S, H, D = 4, 16, 8, 64
        x_template = torch.randn(1, S, H, D)
        x = x_template.expand(B, -1, -1, -1)
        cos = torch.randn(1, S, D // 2).expand(B, -1, -1)
        sin = torch.randn(1, S, D // 2).expand(B, -1, -1)

        output = rope_neox.forward_native(x, cos, sin)

        for i in range(1, B):
            torch.testing.assert_close(output[i], output[0], atol=0, rtol=0)

    # ── Output shape / dtype invariants ───────────────────────────────

    def test_output_shape_preserved(self, rope_neox: RotaryEmbedding) -> None:
        """Output shape must match input shape for all cos/sin dimensionalities."""
        B, S, H, D = 2, 32, 4, 128
        x = torch.randn(B, S, H, D)

        # 2D cos/sin
        cos_2d = torch.randn(S, D // 2)
        sin_2d = torch.randn(S, D // 2)
        assert rope_neox.forward_native(x, cos_2d, sin_2d).shape == (B, S, H, D)

        # 3D cos/sin, matching batch
        cos_3d = torch.randn(B, S, D // 2)
        sin_3d = torch.randn(B, S, D // 2)
        assert rope_neox.forward_native(x, cos_3d, sin_3d).shape == (B, S, H, D)

        # 3D cos/sin, different batch
        cos_3d_1 = torch.randn(1, S, D // 2)
        sin_3d_1 = torch.randn(1, S, D // 2)
        assert rope_neox.forward_native(x, cos_3d_1, sin_3d_1).shape == (B, S, H, D)

    def test_output_dtype_preserved(self, rope_neox: RotaryEmbedding) -> None:
        """Output dtype must match input dtype."""
        for dtype in [torch.float32, torch.float16, torch.bfloat16]:
            x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, dtype=dtype)
            output = rope_neox.forward_native(x, cos, sin)
            assert output.dtype == dtype, f"dtype mismatch: {output.dtype} != {dtype}"

    def test_output_no_nan_inf(
        self,
        rope_neox: RotaryEmbedding,
        rope_gptj: RotaryEmbedding,
    ) -> None:
        """Output must be finite for both RoPE styles."""
        for rope in (rope_neox, rope_gptj):
            x, cos, sin = _make_inputs(batch=2, seq_len=32, num_heads=4, head_dim=64)
            output = rope.forward_native(x, cos, sin)
            assert not torch.isnan(output).any(), f"NaN in {rope}"
            assert not torch.isinf(output).any(), f"Inf in {rope}"

    def test_output_contiguous(self, rope_neox: RotaryEmbedding) -> None:
        """Output should be contiguous."""
        x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64)
        output = rope_neox.forward_native(x, cos, sin)
        assert output.is_contiguous()

    # ── CFG-style doubled batch ───────────────────────────────────────

    def test_cfg_style_doubled_batch(self, rope_neox: RotaryEmbedding) -> None:
        """CFG (classifier-free guidance) doubles batch — cos may have
        original batch size while x has 2× batch."""
        B, S, H, D = 2, 16, 8, 64
        # Simulate CFG: x has 2× batch (cond + uncond concatenated)
        x = torch.randn(B * 2, S, H, D)
        # cos/sin only have the original batch
        cos = torch.randn(B, S, D // 2)
        sin = torch.randn(B, S, D // 2)

        output = rope_neox.forward_native(x, cos, sin)
        assert output.shape == (B * 2, S, H, D)

    # ── RotaryEmbeddingWan forward_native is unchanged ────────────────

    def test_wan_forward_native_unchanged(self) -> None:
        """RotaryEmbeddingWan.forward_native still uses its own implementation
        (unflatten + stack + flatten) and should not crash with 3D cos/sin."""
        from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan

        rope = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True)
        B, S, H, D = 2, 16, 8, 64
        x = torch.randn(B, S, H, D)
        cos = torch.randn(B, S, 1, D // 2)
        sin = torch.randn(B, S, 1, D // 2)

        output = rope.forward_native(x, cos, sin)
        assert output.shape == x.shape
        assert not torch.isnan(output).any()


def test_cuda_paths_fall_back_without_vendored_rotary_kernel() -> None:
    """Source/CPU installs need not contain the wheel-vendored CUDA kernel."""
    x = torch.randn(1, 5, 2, 8)

    rope = RotaryEmbedding(is_neox_style=True)
    rope.apply_rotary_emb_vllm_flash_attn = None
    output = rope.forward_cuda(x, torch.randn(5, 4), torch.randn(5, 4))
    assert output.shape == x.shape

    wan_rope = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True)
    wan_rope.apply_rotary_emb_vllm_flash_attn = None
    output = wan_rope.forward_cuda(
        x,
        torch.randn(1, 5, 1, 4),
        torch.randn(1, 5, 1, 4),
    )
    assert output.shape == x.shape


def test_mindie_full_dim_rope_keeps_h3_cos_sin_layout(monkeypatch: pytest.MonkeyPatch) -> None:
    """H3 passes complete rotary dims to MindIE instead of the CUDA half-dim layout."""
    captured: dict[str, object] = {}

    def rotary_position_embedding(x, cos, sin, **kwargs):
        captured.update(cos=cos, sin=sin, kwargs=kwargs)
        return x

    monkeypatch.setitem(
        sys.modules,
        "mindiesd",
        types.SimpleNamespace(rotary_position_embedding=rotary_position_embedding),
    )
    x = torch.randn(11, 3, 96, dtype=torch.bfloat16)
    cos = torch.randn(11, 96, dtype=torch.bfloat16)
    sin = torch.randn(11, 96, dtype=torch.bfloat16)

    apply_rotary_emb_mindiesd(x, cos, sin, interleaved=False, half_head_dim=False)

    assert captured["cos"] is cos
    assert captured["sin"] is sin
    assert captured["kwargs"] == {
        "rotated_mode": "rotated_half",
        "head_first": False,
        "fused": True,
    }


def test_mindie_rope_restores_3d_layout(monkeypatch: pytest.MonkeyPatch) -> None:
    """mindiesd rotary_position_embedding only accepts 4D ``x``."""
    captured: dict[str, torch.Tensor] = {}

    def rotary_position_embedding(x, cos, sin, **kwargs):
        # Mirror mindiesd's own input guard.
        assert x.dim() == 4, f"mindiesd requires 4D x, got {x.dim()}"
        captured.update(x=x)
        return x

    monkeypatch.setitem(
        sys.modules,
        "mindiesd",
        types.SimpleNamespace(rotary_position_embedding=rotary_position_embedding),
    )

    S, H, D = 11, 3, 96
    x = torch.randn(S, H, D, dtype=torch.bfloat16)
    cos = torch.randn(S, D, dtype=torch.bfloat16)
    sin = torch.randn(S, D, dtype=torch.bfloat16)

    actual = apply_rotary_emb_mindiesd(x, cos, sin, interleaved=False, half_head_dim=False)

    assert captured["x"].shape == (1, S, H, D)
    assert actual.shape == (S, H, D)
    torch.testing.assert_close(actual, x, atol=0, rtol=0)