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

import math

import mlx.core as mx
import mlx.nn as nn

from mlx_music.models.acestep.nn_blocks import (
    MLP,
    TimestepEmbedding,
    timestep_embedding,
)


class TestMLP:
    def test_output_shape(self):
        """MLP preserves [B, T, hidden_size] shape."""
        hidden_size = 64
        intermediate_size = 192
        mlp = MLP(hidden_size, intermediate_size)

        x = mx.random.normal((2, 10, hidden_size))
        out = mlp(x)
        assert out.shape == (
            2,
            10,
            hidden_size,
        ), f"Expected (2, 10, {hidden_size}), got {out.shape}"

    def test_output_shape_single_token(self):
        """MLP works with single-token input."""
        mlp = MLP(hidden_size=128, intermediate_size=384)
        x = mx.random.normal((1, 1, 128))
        out = mlp(x)
        assert out.shape == (1, 1, 128)

    def test_gate_mechanism(self):
        """Verify the SwiGLU gating: down(silu(gate(x)) * up(x))."""
        hidden_size = 32
        intermediate_size = 64
        mlp = MLP(hidden_size, intermediate_size)

        x = mx.random.normal((1, 4, hidden_size))

        # Manual forward pass
        gate_out = nn.silu(mlp.gate_proj(x))
        up_out = mlp.up_proj(x)
        expected = mlp.down_proj(gate_out * up_out)

        actual = mlp(x)
        # Should be numerically identical
        diff = mx.abs(actual - expected).max().item()
        assert diff < 1e-5, f"Gate mechanism mismatch: max diff = {diff}"

    def test_full_config_sizes(self):
        """MLP with ACE-Step production config sizes."""
        mlp = MLP(hidden_size=2048, intermediate_size=6144)
        x = mx.random.normal((1, 8, 2048))
        out = mlp(x)
        assert out.shape == (1, 8, 2048)

    def test_no_bias(self):
        """All projections should have no bias."""
        mlp = MLP(hidden_size=64, intermediate_size=128)
        assert not hasattr(mlp.gate_proj, "bias") or mlp.gate_proj.bias is None
        assert not hasattr(mlp.up_proj, "bias") or mlp.up_proj.bias is None
        assert not hasattr(mlp.down_proj, "bias") or mlp.down_proj.bias is None


class TestTimestepEmbedding:
    def test_sinusoidal_embedding_shape(self):
        """timestep_embedding produces [B, dim] output."""
        t = mx.array([0.1, 0.5, 0.9])
        emb = timestep_embedding(t, dim=256)
        assert emb.shape == (3, 256), f"Expected (3, 256), got {emb.shape}"

    def test_sinusoidal_embedding_odd_dim(self):
        """Odd dim should be zero-padded to correct size."""
        t = mx.array([0.5])
        emb = timestep_embedding(t, dim=255)
        assert emb.shape == (1, 255), f"Expected (1, 255), got {emb.shape}"

    def test_sinusoidal_embedding_different_timesteps(self):
        """Different timesteps should produce different embeddings."""
        t1 = mx.array([0.1])
        t2 = mx.array([0.9])
        emb1 = timestep_embedding(t1, dim=64)
        emb2 = timestep_embedding(t2, dim=64)
        diff = mx.abs(emb1 - emb2).max().item()
        assert diff > 0.01, "Different timesteps should produce different embeddings"

    def test_sinusoidal_embedding_cos_sin_split(self):
        """First half should be cosine, second half sine."""
        t = mx.array([0.5])
        dim = 64
        emb = timestep_embedding(t, dim=dim)

        # Verify manually: first half = cos(args), second half = sin(args)
        half = dim // 2
        freqs = mx.exp(-math.log(10000) * mx.arange(half, dtype=mx.float32) / half)
        args = (t * 1000.0)[:, None] * freqs[None, :]
        expected_cos = mx.cos(args)
        expected_sin = mx.sin(args)

        cos_diff = mx.abs(emb[:, :half] - expected_cos).max().item()
        sin_diff = mx.abs(emb[:, half:] - expected_sin).max().item()
        assert cos_diff < 1e-5, f"Cosine part mismatch: {cos_diff}"
        assert sin_diff < 1e-5, f"Sine part mismatch: {sin_diff}"

    def test_forward_output_shapes(self):
        """TimestepEmbedding.forward returns (temb, timestep_proj) with correct shapes."""
        in_channels = 256
        time_embed_dim = 2048
        te = TimestepEmbedding(in_channels=in_channels, time_embed_dim=time_embed_dim)

        t = mx.array([0.1, 0.5])  # batch of 2
        temb, timestep_proj = te(t)

        assert temb.shape == (
            2,
            time_embed_dim,
        ), f"temb shape: expected (2, {time_embed_dim}), got {temb.shape}"
        assert timestep_proj.shape == (
            2,
            6,
            time_embed_dim,
        ), f"timestep_proj shape: expected (2, 6, {time_embed_dim}), got {timestep_proj.shape}"

    def test_forward_single_sample(self):
        """Single-sample batch should work."""
        te = TimestepEmbedding(in_channels=128, time_embed_dim=512)
        t = mx.array([0.3])
        temb, timestep_proj = te(t)
        assert temb.shape == (1, 512)
        assert timestep_proj.shape == (1, 6, 512)

    def test_timestep_proj_six_vectors(self):
        """timestep_proj should have exactly 6 modulation vectors."""
        te = TimestepEmbedding(in_channels=256, time_embed_dim=1024)
        t = mx.array([0.0, 0.5, 1.0])
        _, timestep_proj = te(t)
        assert (
            timestep_proj.shape[1] == 6
        ), f"Expected 6 modulation vectors, got {timestep_proj.shape[1]}"

    def test_different_timesteps_different_embeddings(self):
        """Different timesteps should produce different temb."""
        te = TimestepEmbedding(in_channels=256, time_embed_dim=512)
        t1 = mx.array([0.1])
        t2 = mx.array([0.9])
        temb1, _ = te(t1)
        temb2, _ = te(t2)
        diff = mx.abs(temb1 - temb2).max().item()
        assert diff > 0.01, "Different timesteps should give different embeddings"

    def test_default_scale(self):
        """Default scale should be 1000."""
        te = TimestepEmbedding(in_channels=256, time_embed_dim=512)
        assert te.scale == 1000.0

    def test_custom_scale(self):
        """Custom scale should be respected."""
        te = TimestepEmbedding(in_channels=256, time_embed_dim=512, scale=500.0)
        assert te.scale == 500.0
