"""Tests for models/acestep/attention.py"""

import mlx.core as mx

from mlx_music.models.acestep.attention import (
    AceStepAttention,
    _make_sliding_window_mask,
    _repeat_kv,
)
from mlx_music.models.acestep.cache import CrossAttentionCache
from mlx_music.models.acestep.config import AceStepDiTConfig


def _small_config(**overrides) -> AceStepDiTConfig:
    """Create a small config for fast tests."""
    defaults = dict(
        hidden_size=64,
        intermediate_size=128,
        num_attention_heads=4,
        num_key_value_heads=2,
        head_dim=16,
        rms_norm_eps=1e-6,
        rope_theta=1000000.0,
        attention_bias=False,
        sliding_window=4,
        num_hidden_layers=4,
        layer_types=[
            "sliding_attention",
            "full_attention",
            "sliding_attention",
            "full_attention",
        ],
    )
    defaults.update(overrides)
    return AceStepDiTConfig(**defaults)


class TestSlidingWindowMask:
    def test_shape(self):
        """Mask should be [1, 1, S, S]."""
        mask = _make_sliding_window_mask(10, window_size=3)
        assert mask.shape == (
            1,
            1,
            10,
            10,
        ), f"Expected (1, 1, 10, 10), got {mask.shape}"

    def test_center_not_masked(self):
        """Diagonal (self-position) should not be masked."""
        mask = _make_sliding_window_mask(8, window_size=2)
        for i in range(8):
            val = mask[0, 0, i, i].item()
            assert val == 0.0, f"Position ({i},{i}) should be 0, got {val}"

    def test_within_window_not_masked(self):
        """Positions within +/- window_size should be 0."""
        window = 3
        S = 10
        mask = _make_sliding_window_mask(S, window_size=window)
        for i in range(S):
            for j in range(S):
                if abs(i - j) <= window:
                    val = mask[0, 0, i, j].item()
                    assert (
                        val == 0.0
                    ), f"Position ({i},{j}) within window should be 0, got {val}"

    def test_outside_window_masked(self):
        """Positions outside window should be -inf."""
        window = 2
        S = 8
        mask = _make_sliding_window_mask(S, window_size=window)
        for i in range(S):
            for j in range(S):
                if abs(i - j) > window:
                    val = mask[0, 0, i, j].item()
                    assert val == float(
                        "-inf"
                    ), f"Position ({i},{j}) outside window should be -inf, got {val}"

    def test_window_larger_than_seq(self):
        """Window larger than sequence = no masking."""
        mask = _make_sliding_window_mask(5, window_size=10)
        # All values should be 0
        assert mx.all(mask == 0.0).item(), "No positions should be masked"


class TestRepeatKV:
    def test_no_repeat(self):
        """num_groups=1 should return input unchanged."""
        x = mx.random.normal((2, 4, 8, 16))
        out = _repeat_kv(x, num_groups=1)
        assert out.shape == x.shape
        assert mx.array_equal(out, x).item()

    def test_repeat_2(self):
        """num_groups=2 should double the head dimension."""
        x = mx.random.normal((2, 4, 8, 16))
        out = _repeat_kv(x, num_groups=2)
        assert out.shape == (2, 8, 8, 16), f"Expected (2, 8, 8, 16), got {out.shape}"

    def test_repeat_content(self):
        """Repeated heads should contain the same values."""
        x = mx.random.normal((1, 2, 4, 8))
        out = _repeat_kv(x, num_groups=3)
        assert out.shape == (1, 6, 4, 8)
        # Head 0 repeated at indices 0,1,2; Head 1 at indices 3,4,5
        for g in range(3):
            assert mx.array_equal(out[:, g, :, :], x[:, 0, :, :]).item()
            assert mx.array_equal(out[:, 3 + g, :, :], x[:, 1, :, :]).item()


class TestAceStepAttentionSelf:
    def test_output_shape(self):
        """Self-attention: [B, T, hidden] -> [B, T, hidden]."""
        config = _small_config()
        attn = AceStepAttention(config, layer_idx=1)  # full_attention layer
        x = mx.random.normal((2, 16, 64))
        out = attn(x)
        assert out.shape == (2, 16, 64), f"Expected (2, 16, 64), got {out.shape}"

    def test_single_token(self):
        """Self-attention works with single-token sequences."""
        config = _small_config()
        attn = AceStepAttention(config, layer_idx=1)
        x = mx.random.normal((1, 1, 64))
        out = attn(x)
        assert out.shape == (1, 1, 64)

    def test_sliding_attention_layer(self):
        """Sliding attention layer (layer_idx=0) should have sliding_window set."""
        config = _small_config()
        attn = AceStepAttention(config, layer_idx=0)  # sliding_attention
        assert attn.sliding_window == 4
        assert attn.attention_type == "sliding_attention"

    def test_full_attention_layer(self):
        """Full attention layer (layer_idx=1) should have no sliding_window."""
        config = _small_config()
        attn = AceStepAttention(config, layer_idx=1)  # full_attention
        assert attn.sliding_window is None
        assert attn.attention_type == "full_attention"

    def test_sliding_attention_output_shape(self):
        """Sliding attention layer still produces correct output shape."""
        config = _small_config()
        attn = AceStepAttention(config, layer_idx=0)  # sliding_attention
        x = mx.random.normal((2, 20, 64))
        out = attn(x)
        assert out.shape == (2, 20, 64)

    def test_gqa_head_counts(self):
        """Verify GQA configuration: 4 Q heads, 2 KV heads, 2 groups."""
        config = _small_config()
        attn = AceStepAttention(config, layer_idx=0)
        assert attn.num_heads == 4
        assert attn.num_kv_heads == 2
        assert attn.num_key_value_groups == 2

    def test_has_rope(self):
        """Self-attention should have RoPE."""
        config = _small_config()
        attn = AceStepAttention(config, layer_idx=0)
        assert hasattr(attn, "rope")

    def test_qk_norm(self):
        """Self-attention should have q_norm and k_norm."""
        config = _small_config()
        attn = AceStepAttention(config, layer_idx=0)
        assert isinstance(attn.q_norm, type(attn.k_norm))

    def test_deterministic_with_same_input(self):
        """Same input should produce same output (no dropout at inference)."""
        config = _small_config()
        attn = AceStepAttention(config, layer_idx=1)
        x = mx.random.normal((1, 8, 64))
        out1 = attn(x)
        out2 = attn(x)
        mx.eval(out1, out2)
        diff = mx.abs(out1 - out2).max().item()
        assert diff < 1e-5, f"Non-deterministic output: max diff = {diff}"


class TestAceStepAttentionCross:
    def test_output_shape(self):
        """Cross-attention: Q from hidden_states, K/V from encoder_hidden_states."""
        config = _small_config()
        attn = AceStepAttention(config, layer_idx=0, is_cross_attention=True)
        hidden = mx.random.normal((2, 16, 64))
        encoder = mx.random.normal((2, 32, 64))
        out = attn(hidden, encoder_hidden_states=encoder)
        assert out.shape == (2, 16, 64), f"Expected (2, 16, 64), got {out.shape}"

    def test_no_rope(self):
        """Cross-attention should NOT have RoPE."""
        config = _small_config()
        attn = AceStepAttention(config, layer_idx=0, is_cross_attention=True)
        assert not hasattr(attn, "rope")

    def test_no_sliding_window(self):
        """Cross-attention should never use sliding window."""
        config = _small_config()
        # Even on a sliding_attention layer, cross-attention has no window
        attn = AceStepAttention(config, layer_idx=0, is_cross_attention=True)
        assert attn.sliding_window is None

    def test_different_encoder_lengths(self):
        """Cross-attention should handle different encoder sequence lengths."""
        config = _small_config()
        attn = AceStepAttention(config, layer_idx=0, is_cross_attention=True)

        hidden = mx.random.normal((1, 8, 64))
        for enc_len in [4, 8, 16, 64]:
            encoder = mx.random.normal((1, enc_len, 64))
            out = attn(hidden, encoder_hidden_states=encoder)
            assert out.shape == (1, 8, 64), f"Failed with encoder length {enc_len}"

    def test_cross_attn_fallback_to_self(self):
        """Cross-attention with encoder_hidden_states=None should use self-attention path."""
        config = _small_config()
        attn = AceStepAttention(config, layer_idx=1, is_cross_attention=True)
        # When encoder_hidden_states is None, should fall back to self-attention
        # (but without RoPE since it's a cross-attention module)
        hidden = mx.random.normal((1, 8, 64))
        out = attn(hidden, encoder_hidden_states=None)
        assert out.shape == (1, 8, 64)


class TestSlidingWindowBehavior:
    def test_sliding_mask_applied(self):
        """Verify sliding window actually restricts attention span.

        With a very small window, distant positions should have near-zero
        influence compared to nearby positions.
        """
        config = _small_config(sliding_window=2)
        attn = AceStepAttention(config, layer_idx=0)  # sliding_attention

        # Create input where first and last tokens are very different
        S = 16
        x = mx.random.normal((1, S, 64))

        # Run self-attention with sliding window
        out = attn(x)
        mx.eval(out)

        # The output should exist and have the right shape
        assert out.shape == (1, S, 64)

    def test_mask_values_128_window(self):
        """Test sliding window mask with the default ACE-Step window of 128."""
        mask = _make_sliding_window_mask(256, window_size=128)
        assert mask.shape == (1, 1, 256, 256)

        # Position (0, 0): distance=0 <= 128 -> should be 0
        assert mask[0, 0, 0, 0].item() == 0.0
        # Position (0, 128): distance=128 <= 128 -> should be 0
        assert mask[0, 0, 0, 128].item() == 0.0
        # Position (0, 129): distance=129 > 128 -> should be -inf
        assert mask[0, 0, 0, 129].item() == float("-inf")
        # Position (0, 255): distance=255 > 128 -> should be -inf
        assert mask[0, 0, 0, 255].item() == float("-inf")


class TestProductionConfig:
    def test_full_config_self_attention(self):
        """Self-attention with production config sizes (hidden=2048, heads=16, kv=8)."""
        config = AceStepDiTConfig()  # Full production config
        attn = AceStepAttention(config, layer_idx=0)
        x = mx.random.normal((1, 8, 2048))
        out = attn(x)
        assert out.shape == (1, 8, 2048)

    def test_full_config_cross_attention(self):
        """Cross-attention with production config sizes."""
        config = AceStepDiTConfig()
        attn = AceStepAttention(config, layer_idx=0, is_cross_attention=True)
        hidden = mx.random.normal((1, 8, 2048))
        encoder = mx.random.normal((1, 16, 2048))
        out = attn(hidden, encoder_hidden_states=encoder)
        assert out.shape == (1, 8, 2048)


class TestCrossAttentionCache:
    """Cross-attention caching produces bit-identical results."""

    def setup_method(self):
        self.config = _small_config()
        self.layer_idx = 0
        self.attn = AceStepAttention(
            self.config, self.layer_idx, is_cross_attention=True
        )
        mx.eval(self.attn.parameters())

    def test_cache_none_same_as_before(self):
        """Passing cache=None should not change output."""
        hidden = mx.random.normal((1, 10, 64))
        encoder = mx.random.normal((1, 20, 64))
        mx.eval(hidden, encoder)

        out_no_cache = self.attn(hidden, encoder_hidden_states=encoder)
        out_none_cache = self.attn(hidden, encoder_hidden_states=encoder, cache=None)
        mx.eval(out_no_cache, out_none_cache)

        assert mx.allclose(out_no_cache, out_none_cache, atol=1e-6)

    def test_cached_output_matches_fresh(self):
        """Cached K/V must produce identical output to fresh computation."""
        hidden = mx.random.normal((1, 10, 64))
        encoder = mx.random.normal((1, 20, 64))
        mx.eval(hidden, encoder)

        # Fresh computation (no cache)
        out_fresh = self.attn(hidden, encoder_hidden_states=encoder)
        mx.eval(out_fresh)

        # First call with cache (populates)
        cache = CrossAttentionCache()
        out_step0 = self.attn(hidden, encoder_hidden_states=encoder, cache=cache)
        mx.eval(out_step0)
        assert cache.has_layer(self.layer_idx)

        # Second call with cache (retrieves) — different hidden_states, same encoder
        hidden2 = mx.random.normal((1, 10, 64))
        mx.eval(hidden2)
        out_cached = self.attn(hidden2, encoder_hidden_states=encoder, cache=cache)
        out_fresh2 = self.attn(hidden2, encoder_hidden_states=encoder, cache=None)
        mx.eval(out_cached, out_fresh2)

        assert mx.allclose(out_cached, out_fresh2, atol=1e-6)

    def test_self_attn_ignores_cache(self):
        """Self-attention should ignore the cache parameter entirely."""
        self_attn = AceStepAttention(self.config, 0, is_cross_attention=False)
        mx.eval(self_attn.parameters())

        hidden = mx.random.normal((1, 10, 64))
        mx.eval(hidden)

        cache = CrossAttentionCache()
        out_no_cache = self_attn(hidden)
        out_with_cache = self_attn(hidden, cache=cache)
        mx.eval(out_no_cache, out_with_cache)

        assert mx.allclose(out_no_cache, out_with_cache, atol=1e-6)
        assert len(cache) == 0  # Self-attention should not populate cache
