"""ACE-Step DiT building blocks for MLX.

Ported from:
  docs/reference/modeling_acestep_v15_turbo.py (TimestepEmbedding, Qwen3MLP)

Contains:
  - MLP: Qwen3-style gated MLP (gate_proj + up_proj -> down_proj)
  - TimestepEmbedding: sinusoidal timestep -> MLP -> 6-dim modulation vectors
  - AdaLN modulation pattern (documented, not a separate class)
"""

import math
from typing import Tuple

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


class MLP(nn.Module):
    """Qwen3-style gated MLP (SwiGLU variant).

    Architecture:
        out = down_proj(silu(gate_proj(x)) * up_proj(x))

    This matches the Qwen3MLP used in the original ACE-Step PyTorch code.
    """

    def __init__(self, hidden_size: int, intermediate_size: int):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)

    def __call__(self, x: mx.array) -> mx.array:
        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))


def timestep_embedding(
    t: mx.array, dim: int, scale: float = 1000.0, max_period: float = 10000.0
) -> mx.array:
    """Create sinusoidal timestep embeddings.

    Matches the PyTorch implementation in TimestepEmbedding.timestep_embedding().

    Args:
        t: 1-D array of shape [B] with timestep values (may be fractional).
        dim: Embedding dimension.
        scale: Multiplier applied to timesteps before embedding.
        max_period: Controls minimum frequency of the embeddings.

    Returns:
        Array of shape [B, dim] with sinusoidal positional embeddings.
    """
    t = t * scale
    half = dim // 2
    # freqs: [half] = exp(-log(max_period) * arange(half) / half)
    freqs = mx.exp(-math.log(max_period) * mx.arange(half, dtype=mx.float32) / half)
    # args: [B, half] = t[:, None] * freqs[None, :]
    args = t[:, None].astype(mx.float32) * freqs[None, :]
    # embedding: [B, dim] = [cos(args), sin(args)]
    embedding = mx.concatenate([mx.cos(args), mx.sin(args)], axis=-1)
    if dim % 2:
        embedding = mx.concatenate(
            [embedding, mx.zeros_like(embedding[:, :1])], axis=-1
        )
    return embedding


class TimestepEmbedding(nn.Module):
    """Timestep embedding for diffusion models.

    Converts scalar timestep values into high-dimensional embeddings:
      1. Sinusoidal positional encoding of timesteps -> [B, in_channels]
      2. Two-layer MLP with SiLU activation -> [B, time_embed_dim]
      3. Project to 6 modulation vectors -> [B, 6, time_embed_dim]

    The 6 modulation vectors are used for AdaLN in DiT layers:
      (shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp)

    Ported from: TimestepEmbedding in modeling_acestep_v15_turbo.py (lines 197-251)
    """

    def __init__(
        self,
        in_channels: int = 256,
        time_embed_dim: int = 2048,
        scale: float = 1000.0,
    ):
        super().__init__()
        self.in_channels = in_channels
        self.scale = scale

        self.linear_1 = nn.Linear(in_channels, time_embed_dim, bias=True)
        self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim, bias=True)
        self.time_proj = nn.Linear(time_embed_dim, time_embed_dim * 6)

    def __call__(self, t: mx.array) -> Tuple[mx.array, mx.array]:
        """Embed timesteps into conditioning vectors.

        Args:
            t: Timestep values of shape [B].

        Returns:
            Tuple of:
              - temb: Timestep embedding of shape [B, time_embed_dim]
              - timestep_proj: 6 modulation vectors of shape [B, 6, time_embed_dim]
        """
        # Sinusoidal embedding: [B] -> [B, in_channels]
        t_freq = timestep_embedding(t, self.in_channels, scale=self.scale)

        # Two-layer MLP: [B, in_channels] -> [B, time_embed_dim]
        temb = nn.silu(self.linear_1(t_freq))
        temb = self.linear_2(temb)

        # Project to 6 modulation vectors: [B, time_embed_dim] -> [B, 6, time_embed_dim]
        # PyTorch: self.time_proj(act2(temb)).unflatten(1, (6, -1))
        proj = self.time_proj(nn.silu(temb))  # [B, time_embed_dim * 6]
        time_embed_dim = proj.shape[-1] // 6
        timestep_proj = proj.reshape(
            proj.shape[0], 6, time_embed_dim
        )  # [B, 6, time_embed_dim]

        return temb, timestep_proj


# ---------------------------------------------------------------------------
# AdaLN Modulation Pattern (not a separate class)
# ---------------------------------------------------------------------------
#
# In ACE-Step DiT layers, adaptive layer norm is implemented via a learnable
# scale_shift_table parameter combined with timestep projections:
#
#   scale_shift_table = Parameter(randn(1, 6, hidden_size) / hidden_size**0.5)
#
# During forward pass:
#   shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = \
#       (scale_shift_table + temb).chunk(6, dim=1)
#
#   # Self-attention with AdaLN:
#   norm_x = rms_norm(x) * (1 + scale_msa) + shift_msa
#   attn_out = self_attn(norm_x)
#   x = x + attn_out * gate_msa
#
#   # MLP with AdaLN:
#   norm_x = rms_norm(x) * (1 + scale_mlp) + shift_mlp
#   mlp_out = mlp(norm_x)
#   x = x + mlp_out * gate_mlp
#
# This pattern will be implemented directly in dit.py's AceStepDiTLayer.
# ---------------------------------------------------------------------------
