"""Tests for the AutoencoderOobleck (audio VAE) MLX port."""

import math

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

from mlx_music.models.acestep.config import AceStepVAEConfig
from mlx_music.models.acestep.vae import (
    AutoencoderOobleck,
    OobleckDecoder,
    OobleckDecoderBlock,
    OobleckEncoder,
    OobleckEncoderBlock,
    OobleckResidualUnit,
    Snake1d,
)


# ---------------------------------------------------------------------------
# Small config for fast tests (full config uses 128-dim, 5 blocks)
# ---------------------------------------------------------------------------
def small_config() -> AceStepVAEConfig:
    """Tiny VAE config for unit tests — fast, low memory."""
    return AceStepVAEConfig(
        encoder_hidden_size=8,
        downsampling_ratios=[2, 4],
        channel_multiples=[1, 2],
        decoder_channels=8,
        decoder_input_channels=4,
        audio_channels=2,
        sampling_rate=48000,
    )


# ---------------------------------------------------------------------------
# Snake1d
# ---------------------------------------------------------------------------
class TestSnake1d:
    def test_output_shape(self):
        """Output shape must match input shape."""
        snake = Snake1d(hidden_dim=16)
        x = mx.random.normal((2, 100, 16))
        y = snake(x)
        assert y.shape == x.shape

    def test_activation_changes_values(self):
        """Snake activation should modify the input (not be identity)."""
        snake = Snake1d(hidden_dim=8)
        x = mx.random.normal((1, 50, 8))
        y = snake(x)
        mx.eval(y)
        # The additive term sin^2(alpha*x)/beta should be nonzero for general x
        diff = mx.abs(y - x).sum().item()
        assert diff > 0.0, "Snake activation should change values"

    def test_learnable_parameters(self):
        """Snake should have alpha and beta as trainable parameters."""
        snake = Snake1d(hidden_dim=4)
        params = snake.parameters()
        assert "alpha" in params
        assert "beta" in params
        assert params["alpha"].shape == (4,)
        assert params["beta"].shape == (4,)


# ---------------------------------------------------------------------------
# OobleckResidualUnit
# ---------------------------------------------------------------------------
class TestOobleckResidualUnit:
    def test_shape_preserved(self):
        """Residual unit output shape should match input (within trimming)."""
        dim = 16
        unit = OobleckResidualUnit(dim, dilation=1)
        x = mx.random.normal((1, 64, dim))
        y = unit(x)
        mx.eval(y)
        # Output time may be trimmed slightly but channels stay the same
        assert y.shape[0] == x.shape[0]
        assert y.shape[2] == dim
        # Time dimension should be close (at most a few frames trimmed)
        assert abs(y.shape[1] - x.shape[1]) <= 6

    def test_dilation_3(self):
        """Dilation=3 should also preserve shape (with correct padding)."""
        dim = 8
        unit = OobleckResidualUnit(dim, dilation=3)
        x = mx.random.normal((1, 128, dim))
        y = unit(x)
        mx.eval(y)
        assert y.shape[0] == 1
        assert y.shape[2] == dim

    def test_dilation_9(self):
        """Dilation=9 should also preserve shape (with correct padding)."""
        dim = 8
        unit = OobleckResidualUnit(dim, dilation=9)
        x = mx.random.normal((1, 128, dim))
        y = unit(x)
        mx.eval(y)
        assert y.shape[0] == 1
        assert y.shape[2] == dim

    def test_residual_connection(self):
        """Output should differ from a zero-weight pass (residual adds input)."""
        dim = 8
        unit = OobleckResidualUnit(dim, dilation=1)
        x = mx.ones((1, 32, dim))
        y = unit(x)
        mx.eval(y)
        # With random weights, output should not be all ones
        assert y.shape[2] == dim


# ---------------------------------------------------------------------------
# OobleckEncoderBlock
# ---------------------------------------------------------------------------
class TestOobleckEncoderBlock:
    def test_downsampling(self):
        """Encoder block should downsample time by stride factor."""
        stride = 4
        block = OobleckEncoderBlock(input_dim=8, output_dim=16, stride=stride)
        T_in = 128
        x = mx.random.normal((1, T_in, 8))
        y = block(x)
        mx.eval(y)
        assert y.shape[0] == 1
        assert y.shape[2] == 16
        # Time should be approximately T_in / stride
        expected_T = T_in // stride
        # Allow for padding differences (should be within 1-2 frames)
        assert (
            abs(y.shape[1] - expected_T) <= 2
        ), f"Expected T~{expected_T}, got {y.shape[1]}"

    def test_channel_change(self):
        """Encoder block should change channels from input_dim to output_dim."""
        block = OobleckEncoderBlock(input_dim=8, output_dim=32, stride=2)
        x = mx.random.normal((1, 64, 8))
        y = block(x)
        mx.eval(y)
        assert y.shape[2] == 32


# ---------------------------------------------------------------------------
# OobleckDecoderBlock
# ---------------------------------------------------------------------------
class TestOobleckDecoderBlock:
    def test_upsampling(self):
        """Decoder block should upsample time by stride factor."""
        stride = 4
        block = OobleckDecoderBlock(input_dim=16, output_dim=8, stride=stride)
        T_in = 32
        x = mx.random.normal((1, T_in, 16))
        y = block(x)
        mx.eval(y)
        assert y.shape[0] == 1
        assert y.shape[2] == 8
        expected_T = T_in * stride
        # ConvTranspose1d may produce slightly different lengths
        assert (
            abs(y.shape[1] - expected_T) <= stride
        ), f"Expected T~{expected_T}, got {y.shape[1]}"

    def test_channel_change(self):
        """Decoder block should change channels from input_dim to output_dim."""
        block = OobleckDecoderBlock(input_dim=32, output_dim=8, stride=2)
        x = mx.random.normal((1, 16, 32))
        y = block(x)
        mx.eval(y)
        assert y.shape[2] == 8


# ---------------------------------------------------------------------------
# OobleckEncoder
# ---------------------------------------------------------------------------
class TestOobleckEncoder:
    def test_output_shape(self):
        """Encoder should produce [B, T/hop_length, encoder_hidden_size]."""
        cfg = small_config()
        encoder = OobleckEncoder(
            encoder_hidden_size=cfg.encoder_hidden_size,
            audio_channels=cfg.audio_channels,
            downsampling_ratios=cfg.downsampling_ratios,
            channel_multiples=cfg.channel_multiples,
        )
        hop = cfg.hop_length  # 2*4 = 8
        T = hop * 10  # 80 samples
        x = mx.random.normal((1, T, cfg.audio_channels))
        y = encoder(x)
        mx.eval(y)
        assert y.shape[0] == 1
        assert y.shape[2] == cfg.encoder_hidden_size
        # Time should be approximately T / hop_length
        expected_T = T // hop
        assert (
            abs(y.shape[1] - expected_T) <= 2
        ), f"Expected T~{expected_T}, got {y.shape[1]}"

    def test_batch_dimension(self):
        """Encoder should handle batch dimension correctly."""
        cfg = small_config()
        encoder = OobleckEncoder(
            encoder_hidden_size=cfg.encoder_hidden_size,
            audio_channels=cfg.audio_channels,
            downsampling_ratios=cfg.downsampling_ratios,
            channel_multiples=cfg.channel_multiples,
        )
        hop = cfg.hop_length
        x = mx.random.normal((3, hop * 5, cfg.audio_channels))
        y = encoder(x)
        mx.eval(y)
        assert y.shape[0] == 3


# ---------------------------------------------------------------------------
# OobleckDecoder
# ---------------------------------------------------------------------------
class TestOobleckDecoder:
    def test_output_shape(self):
        """Decoder should produce [B, T*hop_length, audio_channels]."""
        cfg = small_config()
        decoder = OobleckDecoder(
            channels=cfg.decoder_channels,
            input_channels=cfg.decoder_input_channels,
            audio_channels=cfg.audio_channels,
            upsampling_ratios=cfg.upsampling_ratios,
            channel_multiples=cfg.channel_multiples,
        )
        T_latent = 10
        z = mx.random.normal((1, T_latent, cfg.decoder_input_channels))
        y = decoder(z)
        mx.eval(y)
        assert y.shape[0] == 1
        assert y.shape[2] == cfg.audio_channels
        # Time should be approximately T_latent * hop_length
        hop = cfg.hop_length
        expected_T = T_latent * hop
        # Allow tolerance for conv padding effects
        assert (
            abs(y.shape[1] - expected_T) <= hop
        ), f"Expected T~{expected_T}, got {y.shape[1]}"


# ---------------------------------------------------------------------------
# AutoencoderOobleck (full model)
# ---------------------------------------------------------------------------
class TestAutoencoderOobleck:
    def test_encode_latent_dim(self):
        """Encode should produce latent with correct dimension."""
        cfg = small_config()
        model = AutoencoderOobleck(cfg)
        hop = cfg.hop_length
        T = hop * 10
        x = mx.random.normal((1, T, cfg.audio_channels))
        latent, posterior = model.encode(x, deterministic=True)
        mx.eval(latent)
        assert latent.shape[0] == 1
        # latent_dim = decoder_input_channels / 2? No:
        # encoder_hidden_size=8, split into mean(4) + scale(4), latent_dim=4
        assert latent.shape[2] == cfg.encoder_hidden_size // 2
        # which equals decoder_input_channels
        assert latent.shape[2] == cfg.decoder_input_channels

    def test_decode_audio_shape(self):
        """Decode should produce audio with correct channels."""
        cfg = small_config()
        model = AutoencoderOobleck(cfg)
        T_latent = 10
        z = mx.random.normal((1, T_latent, cfg.decoder_input_channels))
        audio = model.decode(z)
        mx.eval(audio)
        assert audio.shape[0] == 1
        assert audio.shape[2] == cfg.audio_channels

    def test_encode_decode_roundtrip_shapes(self):
        """Encode then decode should approximately restore time dimension."""
        cfg = small_config()
        model = AutoencoderOobleck(cfg)
        hop = cfg.hop_length
        T = hop * 8
        x = mx.random.normal((1, T, cfg.audio_channels))

        latent, _ = model.encode(x, deterministic=True)
        mx.eval(latent)
        reconstructed = model.decode(latent)
        mx.eval(reconstructed)

        assert reconstructed.shape[0] == x.shape[0]
        assert reconstructed.shape[2] == cfg.audio_channels
        # Time should be approximately preserved (within hop_length tolerance)
        assert (
            abs(reconstructed.shape[1] - T) <= hop
        ), f"Expected T~{T}, got {reconstructed.shape[1]}"

    def test_forward_returns_tuple(self):
        """Full forward should return (reconstructed, posterior)."""
        cfg = small_config()
        model = AutoencoderOobleck(cfg)
        hop = cfg.hop_length
        T = hop * 5
        x = mx.random.normal((1, T, cfg.audio_channels))
        result = model(x, deterministic=True)
        assert isinstance(result, tuple)
        assert len(result) == 2
        reconstructed, posterior = result
        mx.eval(reconstructed)
        assert reconstructed.shape[0] == 1
        assert reconstructed.shape[2] == cfg.audio_channels

    def test_stochastic_encode(self):
        """Non-deterministic encode should produce different samples each call."""
        cfg = small_config()
        model = AutoencoderOobleck(cfg)
        hop = cfg.hop_length
        T = hop * 5
        x = mx.random.normal((1, T, cfg.audio_channels))

        latent1, _ = model.encode(x, deterministic=False)
        mx.eval(latent1)
        latent2, _ = model.encode(x, deterministic=False)
        mx.eval(latent2)

        diff = mx.abs(latent1 - latent2).sum().item()
        assert diff > 0.0, "Stochastic sampling should produce different latents"

    def test_deterministic_encode(self):
        """Deterministic encode should produce identical results."""
        cfg = small_config()
        model = AutoencoderOobleck(cfg)
        hop = cfg.hop_length
        T = hop * 5
        x = mx.random.normal((1, T, cfg.audio_channels))
        mx.eval(x)

        latent1, _ = model.encode(x, deterministic=True)
        mx.eval(latent1)
        latent2, _ = model.encode(x, deterministic=True)
        mx.eval(latent2)

        diff = mx.abs(latent1 - latent2).sum().item()
        assert diff == 0.0, "Deterministic encode should be identical"


# ---------------------------------------------------------------------------
# Full-size config smoke test
# ---------------------------------------------------------------------------
class TestFullConfig:
    def test_default_config_dimensions(self):
        """Verify the default config produces correct channel progression."""
        cfg = AceStepVAEConfig()
        assert cfg.hop_length == 1920
        assert cfg.encoder_hidden_size == 128
        # encoder_hidden_size (128) = 2 * decoder_input_channels (64)
        # This means the encoder output splits evenly into mean + scale
        assert cfg.encoder_hidden_size == 2 * cfg.decoder_input_channels

    def test_full_config_instantiation(self):
        """Instantiate full-size model (no forward pass — just check it builds)."""
        cfg = AceStepVAEConfig()
        model = AutoencoderOobleck(cfg)
        # Check the model has encoder and decoder
        assert hasattr(model, "encoder")
        assert hasattr(model, "decoder")
        assert model.hop_length == 1920
