"""Tests for the ACE-Step DiT model."""

import unittest

import mlx.core as mx

from mlx_music.models.acestep.config import AceStepDiTConfig
from mlx_music.models.acestep.dit import (
    AceStepConditionEncoder,
    AceStepConditionGenerationModel,
    AceStepDiTLayer,
    AceStepDiTModel,
    AceStepEncoderLayer,
    AceStepLyricEncoder,
    AceStepTimbreEncoder,
    pack_sequences,
)


def _small_config(**overrides) -> AceStepDiTConfig:
    """Create a tiny config for fast unit tests."""
    defaults = dict(
        hidden_size=64,
        intermediate_size=128,
        num_hidden_layers=2,
        num_attention_heads=4,
        num_key_value_heads=2,
        head_dim=16,
        num_lyric_encoder_hidden_layers=1,
        num_timbre_encoder_hidden_layers=1,
        in_channels=24,  # 8 (acoustic) + 8 (src) + 8 (mask)
        audio_acoustic_hidden_dim=8,
        patch_size=2,
        text_hidden_dim=32,
        timbre_hidden_dim=8,
        sliding_window=4,
        rms_norm_eps=1e-6,
        rope_theta=1000000.0,
        attention_bias=False,
        layer_types=["sliding_attention", "full_attention"],
    )
    defaults.update(overrides)
    return AceStepDiTConfig(**defaults)


class TestEncoderLayer(unittest.TestCase):
    def test_output_shape(self):
        cfg = _small_config()
        layer = AceStepEncoderLayer(cfg, layer_idx=0)
        x = mx.random.normal((2, 8, cfg.hidden_size))
        out = layer(x)
        mx.eval(out)
        self.assertEqual(out.shape, (2, 8, cfg.hidden_size))

    def test_residual_connection(self):
        """Output should differ from input (not identity)."""
        cfg = _small_config()
        layer = AceStepEncoderLayer(cfg, layer_idx=0)
        x = mx.random.normal((1, 4, cfg.hidden_size))
        out = layer(x)
        mx.eval(out)
        self.assertFalse(mx.allclose(out, x).item())

    def test_batch_independence(self):
        """Each batch item should produce independent results."""
        cfg = _small_config()
        layer = AceStepEncoderLayer(cfg, layer_idx=0)
        x = mx.random.normal((2, 4, cfg.hidden_size))
        out = layer(x)
        mx.eval(out)
        out_0 = layer(x[0:1])
        mx.eval(out_0)
        self.assertTrue(mx.allclose(out[0:1], out_0, atol=1e-5).item())


class TestDiTLayer(unittest.TestCase):
    def test_output_shape(self):
        cfg = _small_config()
        layer = AceStepDiTLayer(cfg, layer_idx=0, use_cross_attention=True)
        x = mx.random.normal((2, 8, cfg.hidden_size))
        temb = mx.random.normal((2, 6, cfg.hidden_size))
        enc = mx.random.normal((2, 4, cfg.hidden_size))
        out = layer(x, temb=temb, encoder_hidden_states=enc)
        mx.eval(out)
        self.assertEqual(out.shape, (2, 8, cfg.hidden_size))

    def test_no_cross_attention(self):
        cfg = _small_config()
        layer = AceStepDiTLayer(cfg, layer_idx=0, use_cross_attention=False)
        x = mx.random.normal((1, 6, cfg.hidden_size))
        temb = mx.random.normal((1, 6, cfg.hidden_size))
        out = layer(x, temb=temb)
        mx.eval(out)
        self.assertEqual(out.shape, (1, 6, cfg.hidden_size))

    def test_adaln_modulation(self):
        """Different timestep embeddings should produce different outputs."""
        cfg = _small_config()
        layer = AceStepDiTLayer(cfg, layer_idx=0, use_cross_attention=False)
        x = mx.random.normal((1, 4, cfg.hidden_size))
        temb_a = mx.zeros((1, 6, cfg.hidden_size))
        temb_b = mx.ones((1, 6, cfg.hidden_size))
        out_a = layer(x, temb=temb_a)
        out_b = layer(x, temb=temb_b)
        mx.eval(out_a, out_b)
        self.assertFalse(mx.allclose(out_a, out_b).item())


class TestDiTModel(unittest.TestCase):
    def test_output_shape(self):
        cfg = _small_config()
        model = AceStepDiTModel(cfg)
        B, T = 1, 8
        acoustic = mx.random.normal((B, T, cfg.audio_acoustic_hidden_dim))
        context = mx.random.normal(
            (B, T, cfg.in_channels - cfg.audio_acoustic_hidden_dim)
        )
        timestep = mx.array([0.5])
        timestep_r = mx.array([0.5])
        enc = mx.random.normal((B, 4, cfg.hidden_size))

        out = model(acoustic, timestep, timestep_r, enc, context)
        mx.eval(out)
        self.assertEqual(out.shape, (B, T, cfg.audio_acoustic_hidden_dim))

    def test_odd_seq_len_padding(self):
        """Sequence not divisible by patch_size should still work."""
        cfg = _small_config(patch_size=4)
        model = AceStepDiTModel(cfg)
        B, T = 1, 10  # 10 % 4 != 0
        acoustic = mx.random.normal((B, T, cfg.audio_acoustic_hidden_dim))
        context = mx.random.normal(
            (B, T, cfg.in_channels - cfg.audio_acoustic_hidden_dim)
        )
        timestep = mx.array([0.8])
        enc = mx.random.normal((B, 2, cfg.hidden_size))

        out = model(acoustic, timestep, timestep, enc, context)
        mx.eval(out)
        # Output must match original T, not padded
        self.assertEqual(out.shape[1], T)

    def test_different_timesteps_differ(self):
        """Different timesteps should produce different outputs."""
        cfg = _small_config()
        model = AceStepDiTModel(cfg)
        B, T = 1, 4
        acoustic = mx.random.normal((B, T, cfg.audio_acoustic_hidden_dim))
        context = mx.random.normal(
            (B, T, cfg.in_channels - cfg.audio_acoustic_hidden_dim)
        )
        enc = mx.random.normal((B, 2, cfg.hidden_size))

        out_a = model(acoustic, mx.array([0.1]), mx.array([0.1]), enc, context)
        out_b = model(acoustic, mx.array([0.9]), mx.array([0.9]), enc, context)
        mx.eval(out_a, out_b)
        self.assertFalse(mx.allclose(out_a, out_b).item())


class TestLyricEncoder(unittest.TestCase):
    def test_output_shape(self):
        cfg = _small_config()
        encoder = AceStepLyricEncoder(cfg)
        x = mx.random.normal((1, 12, cfg.text_hidden_dim))
        mask = mx.ones((1, 12))
        out = encoder(x, mask)
        mx.eval(out)
        self.assertEqual(out.shape, (1, 12, cfg.hidden_size))

    def test_projects_to_hidden_size(self):
        """Input text_hidden_dim should be projected to hidden_size."""
        cfg = _small_config(text_hidden_dim=32, hidden_size=64)
        encoder = AceStepLyricEncoder(cfg)
        x = mx.random.normal((1, 4, 32))
        out = encoder(x)
        mx.eval(out)
        self.assertEqual(out.shape[-1], 64)


class TestTimbreEncoder(unittest.TestCase):
    def test_output_shape(self):
        cfg = _small_config()
        encoder = AceStepTimbreEncoder(cfg)
        x = mx.random.normal((2, 10, cfg.timbre_hidden_dim))
        embs, mask = encoder(x)
        mx.eval(embs, mask)
        self.assertEqual(embs.shape, (2, 1, cfg.hidden_size))
        self.assertEqual(mask.shape, (2, 1))


class TestPackSequences(unittest.TestCase):
    def test_all_valid(self):
        """When all masks are 1, output should be simple concatenation."""
        h1 = mx.ones((1, 3, 4))
        h2 = mx.ones((1, 2, 4)) * 2
        m1 = mx.ones((1, 3))
        m2 = mx.ones((1, 2))
        packed, mask = pack_sequences(h1, h2, m1, m2)
        mx.eval(packed, mask)
        self.assertEqual(packed.shape, (1, 5, 4))
        self.assertEqual(mask.shape, (1, 5))
        # All positions should be valid
        self.assertTrue(mx.all(mask == 1.0).item())

    def test_with_padding(self):
        """Padded tokens (mask=0) should be sorted to the end."""
        h1 = mx.ones((1, 2, 4))
        h2 = mx.ones((1, 2, 4)) * 2
        m1 = mx.array([[1.0, 0.0]])  # second is padding
        m2 = mx.array([[1.0, 1.0]])  # both valid
        packed, mask = pack_sequences(h1, h2, m1, m2)
        mx.eval(packed, mask)
        # 3 valid tokens, 1 padding
        self.assertEqual(mask.sum().item(), 3.0)
        self.assertEqual(packed.shape, (1, 4, 4))


class TestConditionEncoder(unittest.TestCase):
    def test_output_shape(self):
        cfg = _small_config()
        encoder = AceStepConditionEncoder(cfg)
        B = 1
        text = mx.random.normal((B, 5, cfg.text_hidden_dim))
        text_mask = mx.ones((B, 5))
        lyrics = mx.random.normal((B, 8, cfg.text_hidden_dim))
        lyric_mask = mx.ones((B, 8))

        out, mask = encoder(text, text_mask, lyrics, lyric_mask)
        mx.eval(out, mask)
        # Packed length: lyrics(8) + timbre(1, placeholder) + text(5) = 14
        self.assertEqual(out.shape[0], B)
        self.assertEqual(out.shape[2], cfg.hidden_size)
        self.assertEqual(out.shape[1], mask.shape[1])

    def test_with_timbre(self):
        cfg = _small_config()
        encoder = AceStepConditionEncoder(cfg)
        B = 1
        text = mx.random.normal((B, 3, cfg.text_hidden_dim))
        text_mask = mx.ones((B, 3))
        lyrics = mx.random.normal((B, 4, cfg.text_hidden_dim))
        lyric_mask = mx.ones((B, 4))
        timbre = mx.random.normal((B, 6, cfg.timbre_hidden_dim))

        out, mask = encoder(text, text_mask, lyrics, lyric_mask, timbre)
        mx.eval(out, mask)
        self.assertEqual(out.shape[0], B)
        self.assertEqual(out.shape[2], cfg.hidden_size)


class TestConditionGenerationModel(unittest.TestCase):
    def test_denoise_step(self):
        """Full forward pass through the generation model."""
        cfg = _small_config()
        model = AceStepConditionGenerationModel(cfg)
        B, T = 1, 4

        # Prepare conditions
        text = mx.random.normal((B, 3, cfg.text_hidden_dim))
        text_mask = mx.ones((B, 3))
        lyrics = mx.random.normal((B, 4, cfg.text_hidden_dim))
        lyric_mask = mx.ones((B, 4))
        context = mx.random.normal(
            (B, T, cfg.in_channels - cfg.audio_acoustic_hidden_dim)
        )

        enc_h, enc_m, ctx = model.prepare_condition(
            text, text_mask, lyrics, lyric_mask, context
        )
        mx.eval(enc_h, enc_m, ctx)

        # Denoise
        xt = mx.random.normal((B, T, cfg.audio_acoustic_hidden_dim))
        vt = model.denoise_step(xt, 0.5, enc_h, ctx)
        mx.eval(vt)
        self.assertEqual(vt.shape, (B, T, cfg.audio_acoustic_hidden_dim))

    def test_output_not_nan(self):
        cfg = _small_config()
        model = AceStepConditionGenerationModel(cfg)
        B, T = 1, 4

        text = mx.random.normal((B, 2, cfg.text_hidden_dim))
        text_mask = mx.ones((B, 2))
        lyrics = mx.random.normal((B, 2, cfg.text_hidden_dim))
        lyric_mask = mx.ones((B, 2))
        context = mx.random.normal(
            (B, T, cfg.in_channels - cfg.audio_acoustic_hidden_dim)
        )

        enc_h, enc_m, ctx = model.prepare_condition(
            text, text_mask, lyrics, lyric_mask, context
        )
        xt = mx.random.normal((B, T, cfg.audio_acoustic_hidden_dim))
        vt = model.denoise_step(xt, 0.5, enc_h, ctx)
        mx.eval(vt)
        self.assertFalse(mx.any(mx.isnan(vt)).item())


class TestDiTCrossAttentionCache(unittest.TestCase):
    """End-to-end cache test: denoise_step calls with cache produce
    identical output to calls without cache."""

    def setUp(self):
        from mlx_music.models.acestep.cache import CrossAttentionCache

        self.CrossAttentionCache = CrossAttentionCache
        self.config = _small_config()
        self.model = AceStepConditionGenerationModel(self.config)
        mx.eval(self.model.parameters())

    def test_cached_denoise_step_matches_fresh(self):
        B, T = 1, 4
        enc_len = 6
        acoustic_dim = self.config.audio_acoustic_hidden_dim

        xt = mx.random.normal((B, T, acoustic_dim))
        encoder_hs = mx.random.normal((B, enc_len, self.config.hidden_size))
        context = mx.random.normal((B, T, self.config.in_channels - acoustic_dim))
        mx.eval(xt, encoder_hs, context)

        # Without cache
        out1_no_cache = self.model.denoise_step(xt, 1.0, encoder_hs, context)
        out2_no_cache = self.model.denoise_step(xt, 0.9, encoder_hs, context)
        mx.eval(out1_no_cache, out2_no_cache)

        # With cache
        cache = self.CrossAttentionCache()
        out1_cached = self.model.denoise_step(xt, 1.0, encoder_hs, context, cache=cache)
        self.assertEqual(len(cache), self.config.num_hidden_layers)

        out2_cached = self.model.denoise_step(xt, 0.9, encoder_hs, context, cache=cache)
        mx.eval(out1_cached, out2_cached)

        self.assertTrue(mx.allclose(out1_no_cache, out1_cached, atol=1e-5).item())
        self.assertTrue(mx.allclose(out2_no_cache, out2_cached, atol=1e-5).item())


class TestDebugCapture(unittest.TestCase):
    """Test that debug_capture records intermediates correctly."""

    def test_captures_intermediates(self):
        from mlx_music.debug.hooks import IntermediateCapture

        cfg = _small_config()
        model = AceStepConditionGenerationModel(cfg)
        mx.eval(model.parameters())

        B, T = 1, 4
        xt = mx.random.normal((B, T, cfg.audio_acoustic_hidden_dim))
        enc = mx.random.normal((B, 6, cfg.hidden_size))
        ctx = mx.random.normal((B, T, cfg.in_channels - cfg.audio_acoustic_hidden_dim))
        mx.eval(xt, enc, ctx)

        cap = IntermediateCapture()
        out = model.denoise_step(xt, 0.5, enc, ctx, debug_capture=cap)
        mx.eval(out)

        # Should have: post_patch_embed, encoder_projected, timestep_proj,
        # post_layer_0, post_layer_1, post_adaln_out, final_output
        self.assertIn("post_patch_embed", cap)
        self.assertIn("encoder_projected", cap)
        self.assertIn("timestep_proj", cap)
        self.assertIn("post_layer_0", cap)
        self.assertIn(f"post_layer_{cfg.num_hidden_layers - 1}", cap)
        self.assertIn("post_adaln_out", cap)
        self.assertIn("final_output", cap)
        expected_keys = 3 + cfg.num_hidden_layers + 2  # 3 initial + layers + 2 output
        self.assertEqual(len(cap), expected_keys)


if __name__ == "__main__":
    unittest.main()
