"""ACE-Step DiT (Diffusion Transformer) for MLX — inference only.

Ported from:
  docs/reference/modeling_acestep_v15_turbo.py

Contains:
  - AceStepEncoderLayer: bidirectional self-attn + MLP (lyric/timbre encoders)
  - AceStepDiTLayer: AdaLN self-attn + cross-attn + AdaLN MLP
  - AceStepDiTModel: patch_in → N DiT layers → AdaLN norm → patch_out
  - AceStepLyricEncoder: text projection + encoder layers
  - AceStepTimbreEncoder: acoustic projection + encoder layers + CLS pooling
  - AceStepConditionEncoder: text + lyrics + timbre → packed sequence
  - AceStepConditionGenerationModel: top-level wrapper (condition encoder + DiT)
"""

from __future__ import annotations

import math
from typing import TYPE_CHECKING, Optional, Tuple

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

from .attention import AceStepAttention, _make_sliding_window_mask
from .config import AceStepDiTConfig
from .nn_blocks import MLP, TimestepEmbedding

if TYPE_CHECKING:
    from .cache import CrossAttentionCache
    from ...debug.hooks import IntermediateCapture


# ---------------------------------------------------------------------------
# Encoder layer (bidirectional self-attention + MLP, used by lyric/timbre)
# ---------------------------------------------------------------------------


class AceStepEncoderLayer(nn.Module):
    """Bidirectional encoder layer: RMSNorm → self-attn + residual → RMSNorm → MLP + residual."""

    def __init__(self, config: AceStepDiTConfig, layer_idx: int):
        super().__init__()
        self.self_attn = AceStepAttention(config, layer_idx, is_cross_attention=False)
        self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = nn.RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
        self.mlp = MLP(config.hidden_size, config.intermediate_size)
        self.attention_type = config.layer_types[layer_idx]

    def __call__(
        self,
        hidden_states: mx.array,
        mask: Optional[mx.array] = None,
    ) -> mx.array:
        # Self-attention with residual
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states = self.self_attn(hidden_states, mask=mask)
        hidden_states = residual + hidden_states

        # MLP with residual
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        return hidden_states


# ---------------------------------------------------------------------------
# DiT layer (AdaLN self-attn + cross-attn + AdaLN MLP)
# ---------------------------------------------------------------------------


class AceStepDiTLayer(nn.Module):
    """DiT layer with AdaLN modulation from timestep embeddings.

    Forward flow:
      1. AdaLN self-attention: norm(x)*(1+scale)+shift → attn → x + out*gate
      2. Cross-attention: norm(x) → cross-attn(encoder_hidden_states) → x + out
      3. AdaLN MLP: norm(x)*(1+scale)+shift → MLP → x + out*gate

    The 6 modulation vectors come from scale_shift_table + timestep_proj:
      (shift_sa, scale_sa, gate_sa, shift_mlp, scale_mlp, gate_mlp)
    """

    def __init__(
        self,
        config: AceStepDiTConfig,
        layer_idx: int,
        use_cross_attention: bool = True,
    ):
        super().__init__()
        self.attention_type = config.layer_types[layer_idx]

        # Self-attention
        self.self_attn_norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.self_attn = AceStepAttention(config, layer_idx, is_cross_attention=False)

        # Cross-attention
        self.use_cross_attention = use_cross_attention
        if use_cross_attention:
            self.cross_attn_norm = nn.RMSNorm(
                config.hidden_size, eps=config.rms_norm_eps
            )
            self.cross_attn = AceStepAttention(
                config, layer_idx, is_cross_attention=True
            )

        # MLP
        self.mlp_norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.mlp = MLP(config.hidden_size, config.intermediate_size)

        # AdaLN scale-shift-gate table: [1, 6, hidden_size]
        self.scale_shift_table = mx.random.normal((1, 6, config.hidden_size)) * (
            config.hidden_size**-0.5
        )

    def __call__(
        self,
        hidden_states: mx.array,
        temb: mx.array,
        mask: Optional[mx.array] = None,
        encoder_hidden_states: Optional[mx.array] = None,
        encoder_mask: Optional[mx.array] = None,
        cache: Optional[CrossAttentionCache] = None,
    ) -> mx.array:
        """
        Args:
            hidden_states: [B, S, D]
            temb: timestep projection [B, 6, D] (from TimestepEmbedding)
            mask: self-attention mask [1, 1, S, S] or None
            encoder_hidden_states: [B, S_enc, D] for cross-attention
            encoder_mask: [1, 1, S_q, S_enc] or None
            cache: Optional cross-attention KV cache
        """
        # Unpack 6 modulation vectors: each [B, 1, D]
        mod = self.scale_shift_table + temb  # [B, 6, D]
        shift_sa = mod[:, 0:1, :]
        scale_sa = mod[:, 1:2, :]
        gate_sa = mod[:, 2:3, :]
        shift_mlp = mod[:, 3:4, :]
        scale_mlp = mod[:, 4:5, :]
        gate_mlp = mod[:, 5:6, :]

        # 1. AdaLN self-attention
        normed = self.self_attn_norm(hidden_states) * (1.0 + scale_sa) + shift_sa
        attn_out = self.self_attn(normed, mask=mask)
        hidden_states = hidden_states + attn_out * gate_sa

        # 2. Cross-attention (no AdaLN — just norm + residual)
        if self.use_cross_attention and encoder_hidden_states is not None:
            normed = self.cross_attn_norm(hidden_states)
            cross_out = self.cross_attn(
                normed,
                encoder_hidden_states=encoder_hidden_states,
                mask=encoder_mask,
                cache=cache,
            )
            hidden_states = hidden_states + cross_out

        # 3. AdaLN MLP
        normed = self.mlp_norm(hidden_states) * (1.0 + scale_mlp) + shift_mlp
        ff_out = self.mlp(normed)
        hidden_states = hidden_states + ff_out * gate_mlp

        return hidden_states


# ---------------------------------------------------------------------------
# DiT model (patch embed → layers → AdaLN → patch unembed)
# ---------------------------------------------------------------------------


class AceStepDiTModel(nn.Module):
    """DiT backbone: converts audio latent patches through N transformer layers.

    Input:  [B, T, in_channels]  (acoustic + context latents concatenated)
    Output: [B, T, audio_acoustic_hidden_dim]  (predicted velocity)

    Patch embedding via Conv1d stride=patch_size, patch unembedding via
    ConvTranspose1d. Timestep conditioning via two TimestepEmbedding modules
    (one for t, one for t-r).
    """

    def __init__(self, config: AceStepDiTConfig):
        super().__init__()
        self.config = config
        inner_dim = config.hidden_size
        self.patch_size = config.patch_size

        # Patch embed: [B, T, C] → Conv1d → [B, T//patch, D]
        self.proj_in = nn.Conv1d(
            config.in_channels,
            inner_dim,
            kernel_size=config.patch_size,
            stride=config.patch_size,
            padding=0,
        )

        # Timestep embeddings
        self.time_embed = TimestepEmbedding(in_channels=256, time_embed_dim=inner_dim)
        self.time_embed_r = TimestepEmbedding(in_channels=256, time_embed_dim=inner_dim)

        # Condition projector
        self.condition_embedder = nn.Linear(inner_dim, inner_dim, bias=True)

        # RoPE
        self.rope = nn.RoPE(config.head_dim, traditional=False, base=config.rope_theta)

        # DiT layers
        self.layers = [
            AceStepDiTLayer(config, i, use_cross_attention=True)
            for i in range(config.num_hidden_layers)
        ]

        # Output: AdaLN + ConvTranspose1d
        self.norm_out = nn.RMSNorm(inner_dim, eps=config.rms_norm_eps)
        self.proj_out = nn.ConvTranspose1d(
            inner_dim,
            config.audio_acoustic_hidden_dim,
            kernel_size=config.patch_size,
            stride=config.patch_size,
            padding=0,
        )
        # Output AdaLN: 2 modulation vectors (shift, scale)
        self.scale_shift_table = mx.random.normal((1, 2, inner_dim)) * (inner_dim**-0.5)

    def __call__(
        self,
        hidden_states: mx.array,
        timestep: mx.array,
        timestep_r: mx.array,
        encoder_hidden_states: mx.array,
        context_latents: mx.array,
        cache: Optional[CrossAttentionCache] = None,
        debug_capture: Optional[IntermediateCapture] = None,
    ) -> mx.array:
        """
        Args:
            hidden_states: Audio latent [B, T, acoustic_dim] (64-dim)
            timestep: [B] float timestep values
            timestep_r: [B] float (for turbo inference, typically r=t)
            encoder_hidden_states: [B, S_enc, D] packed condition sequence
            context_latents: [B, T, 128] (src_latent 64 + chunk_mask 64)
            cache: Optional cross-attention KV cache
            debug_capture: Optional intermediate tensor capture for debugging

        Returns:
            Predicted velocity [B, T, acoustic_dim]
        """
        # Concatenate context with hidden states: [B, T, in_channels]
        hidden_states = mx.concatenate([context_latents, hidden_states], axis=-1)

        original_seq_len = hidden_states.shape[1]

        # Pad to multiple of patch_size
        pad_length = 0
        if hidden_states.shape[1] % self.patch_size != 0:
            pad_length = self.patch_size - (hidden_states.shape[1] % self.patch_size)
            hidden_states = mx.pad(
                hidden_states,
                [(0, 0), (0, pad_length), (0, 0)],
            )

        # Patch embed: Conv1d expects [B, T, C] in MLX (channels-last)
        hidden_states = self.proj_in(hidden_states)  # [B, T//patch, D]

        # Project encoder states
        encoder_hidden_states = self.condition_embedder(encoder_hidden_states)

        # Compute timestep embeddings
        temb_t, timestep_proj_t = self.time_embed(timestep)
        temb_r, timestep_proj_r = self.time_embed_r(timestep - timestep_r)
        temb = temb_t + temb_r  # [B, D] — used for output AdaLN
        timestep_proj = timestep_proj_t + timestep_proj_r  # [B, 6, D]

        if debug_capture is not None:
            debug_capture.record("post_patch_embed", hidden_states)
            debug_capture.record("encoder_projected", encoder_hidden_states)
            debug_capture.record("timestep_proj", timestep_proj)

        # Build self-attention masks
        seq_len = hidden_states.shape[1]
        sliding_mask = _make_sliding_window_mask(
            seq_len, self.config.sliding_window, dtype=hidden_states.dtype
        )
        # Full attention = no mask (None). Cross-attention = no mask for now.

        # Pass through DiT layers
        for i, layer in enumerate(self.layers):
            attn_mask = (
                sliding_mask if layer.attention_type == "sliding_attention" else None
            )
            hidden_states = layer(
                hidden_states,
                temb=timestep_proj,
                mask=attn_mask,
                encoder_hidden_states=encoder_hidden_states,
                encoder_mask=None,
                cache=cache,
            )
            if debug_capture is not None:
                debug_capture.record(f"post_layer_{i}", hidden_states)

        # Output AdaLN: norm → scale/shift → ConvTranspose1d
        shift, scale = mx.split(
            self.scale_shift_table + mx.expand_dims(temb, axis=1),  # [B, 2, D]
            2,
            axis=1,
        )
        hidden_states = self.norm_out(hidden_states) * (1.0 + scale) + shift

        if debug_capture is not None:
            debug_capture.record("post_adaln_out", hidden_states)

        # Patch unembed
        hidden_states = self.proj_out(hidden_states)  # [B, T, acoustic_dim]

        # Crop to original length
        hidden_states = hidden_states[:, :original_seq_len, :]

        if debug_capture is not None:
            debug_capture.record("final_output", hidden_states)

        return hidden_states


# ---------------------------------------------------------------------------
# Lyric encoder
# ---------------------------------------------------------------------------


class AceStepLyricEncoder(nn.Module):
    """Encodes lyric hidden states through bidirectional transformer layers.

    Input:  lyric_hidden_states [B, S, text_hidden_dim] from text encoder
    Output: encoded lyrics [B, S, hidden_size]
    """

    def __init__(self, config: AceStepDiTConfig):
        super().__init__()
        self.config = config
        self.embed_tokens = nn.Linear(config.text_hidden_dim, config.hidden_size)
        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.layers = [
            AceStepEncoderLayer(config, i)
            for i in range(config.num_lyric_encoder_hidden_layers)
        ]

    def __call__(
        self,
        inputs_embeds: mx.array,
        attention_mask: Optional[mx.array] = None,
    ) -> mx.array:
        """
        Args:
            inputs_embeds: [B, S, text_hidden_dim] lyric hidden states
            attention_mask: [B, S] (1=valid, 0=pad) — currently unused (no padding)

        Returns:
            [B, S, hidden_size]
        """
        hidden_states = self.embed_tokens(inputs_embeds)
        seq_len = hidden_states.shape[1]

        # Build mask mapping for sliding/full attention layers
        sliding_mask = _make_sliding_window_mask(
            seq_len, self.config.sliding_window, dtype=hidden_states.dtype
        )

        for layer in self.layers:
            mask = sliding_mask if layer.attention_type == "sliding_attention" else None
            hidden_states = layer(hidden_states, mask=mask)

        hidden_states = self.norm(hidden_states)
        return hidden_states


# ---------------------------------------------------------------------------
# Timbre encoder (simplified for inference — no packed batch support)
# ---------------------------------------------------------------------------


class AceStepTimbreEncoder(nn.Module):
    """Encodes reference audio acoustic features into timbre embeddings.

    Simplified for inference: takes [B, T, timbre_hidden_dim] and returns
    [B, 1, hidden_size] (CLS token output). In training, the PyTorch version
    supports packed batches with refer_audio_order_mask — we simplify this
    for single-batch inference.
    """

    def __init__(self, config: AceStepDiTConfig):
        super().__init__()
        self.config = config
        self.embed_tokens = nn.Linear(config.timbre_hidden_dim, config.hidden_size)
        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.special_token = mx.random.normal((1, 1, config.hidden_size)) * 0.02
        self.layers = [
            AceStepEncoderLayer(config, i)
            for i in range(config.num_timbre_encoder_hidden_layers)
        ]

    def __call__(
        self,
        inputs_embeds: mx.array,
        attention_mask: Optional[mx.array] = None,
    ) -> Tuple[mx.array, mx.array]:
        """
        Args:
            inputs_embeds: [B, T, timbre_hidden_dim]

        Returns:
            (timbre_embs, mask): ([B, 1, hidden_size], [B, 1])
        """
        B = inputs_embeds.shape[0]
        hidden_states = self.embed_tokens(inputs_embeds)
        seq_len = hidden_states.shape[1]

        sliding_mask = _make_sliding_window_mask(
            seq_len, self.config.sliding_window, dtype=hidden_states.dtype
        )

        for layer in self.layers:
            mask = sliding_mask if layer.attention_type == "sliding_attention" else None
            hidden_states = layer(hidden_states, mask=mask)

        hidden_states = self.norm(hidden_states)

        # CLS token: take first position as timbre embedding
        cls_output = hidden_states[:, 0:1, :]  # [B, 1, D]
        mask = mx.ones((B, 1))
        return cls_output, mask


# ---------------------------------------------------------------------------
# Sequence packing utility
# ---------------------------------------------------------------------------


def pack_sequences(
    hidden1: mx.array,
    hidden2: mx.array,
    mask1: mx.array,
    mask2: mx.array,
) -> Tuple[mx.array, mx.array]:
    """Pack two sequences by concatenating and sorting valid tokens first.

    In the PyTorch reference, this sorts by mask so valid tokens (mask=1)
    come before padding (mask=0). For inference with no padding, this
    simplifies to concatenation.

    Args:
        hidden1: [B, L1, D]
        hidden2: [B, L2, D]
        mask1: [B, L1]
        mask2: [B, L2]

    Returns:
        (packed_hidden, packed_mask): [B, L1+L2, D], [B, L1+L2]
    """
    packed_hidden = mx.concatenate([hidden1, hidden2], axis=1)
    packed_mask = mx.concatenate([mask1, mask2], axis=1)

    # Sort by mask (valid first) — use argsort descending
    B, L, D = packed_hidden.shape
    # For each batch, sort indices so mask=1 comes before mask=0
    sort_idx = mx.argsort(-packed_mask, axis=1)  # descending by negating

    # Gather hidden states by sorted indices
    # Expand sort_idx to [B, L, D] for gathering
    sort_idx_expanded = mx.repeat(mx.expand_dims(sort_idx, axis=-1), D, axis=-1)
    packed_hidden = mx.take_along_axis(packed_hidden, sort_idx_expanded, axis=1)

    # Recompute mask: first `count` positions are valid
    lengths = packed_mask.sum(axis=1, keepdims=True)  # [B, 1]
    positions = mx.arange(L)[None, :]  # [1, L]
    new_mask = (positions < lengths).astype(mx.float32)

    return packed_hidden, new_mask


# ---------------------------------------------------------------------------
# Condition encoder
# ---------------------------------------------------------------------------


class AceStepConditionEncoder(nn.Module):
    """Encodes text + lyrics + timbre → single packed conditioning sequence.

    Packing order: lyrics + timbre → + text (sorted by mask each time).
    """

    def __init__(self, config: AceStepDiTConfig):
        super().__init__()
        self.text_projector = nn.Linear(
            config.text_hidden_dim, config.hidden_size, bias=False
        )
        self.lyric_encoder = AceStepLyricEncoder(config)
        self.timbre_encoder = AceStepTimbreEncoder(config)

    def __call__(
        self,
        text_hidden_states: mx.array,
        text_attention_mask: mx.array,
        lyric_hidden_states: mx.array,
        lyric_attention_mask: mx.array,
        timbre_hidden_states: Optional[mx.array] = None,
    ) -> Tuple[mx.array, mx.array]:
        """
        Args:
            text_hidden_states: [B, S_text, text_hidden_dim]
            text_attention_mask: [B, S_text]
            lyric_hidden_states: [B, S_lyric, text_hidden_dim]
            lyric_attention_mask: [B, S_lyric]
            timbre_hidden_states: [B, T, timbre_hidden_dim] or None

        Returns:
            (encoder_hidden_states, encoder_attention_mask):
                [B, S_packed, hidden_size], [B, S_packed]
        """
        # Project text
        text_hidden_states = self.text_projector(text_hidden_states)

        # Encode lyrics
        lyric_hidden_states = self.lyric_encoder(
            lyric_hidden_states, lyric_attention_mask
        )

        # Encode timbre (or use empty placeholder)
        if timbre_hidden_states is not None:
            timbre_embs, timbre_mask = self.timbre_encoder(timbre_hidden_states)
        else:
            B = text_hidden_states.shape[0]
            D = self.text_projector.weight.shape[0]
            timbre_embs = mx.zeros((B, 1, D))
            timbre_mask = mx.zeros((B, 1))

        # Pack: lyrics + timbre → + text
        enc_hidden, enc_mask = pack_sequences(
            lyric_hidden_states, timbre_embs, lyric_attention_mask, timbre_mask
        )
        enc_hidden, enc_mask = pack_sequences(
            enc_hidden, text_hidden_states, enc_mask, text_attention_mask
        )
        return enc_hidden, enc_mask


# ---------------------------------------------------------------------------
# Top-level generation model
# ---------------------------------------------------------------------------


class AceStepConditionGenerationModel(nn.Module):
    """Top-level model: condition encoder + DiT decoder.

    For inference: prepare_condition → denoise_step (called by sampler).
    """

    def __init__(self, config: AceStepDiTConfig):
        super().__init__()
        self.config = config
        self.decoder = AceStepDiTModel(config)
        self.encoder = AceStepConditionEncoder(config)
        # Null condition for CFG
        self.null_condition_emb = mx.random.normal((1, 1, config.hidden_size))

    def prepare_condition(
        self,
        text_hidden_states: mx.array,
        text_attention_mask: mx.array,
        lyric_hidden_states: mx.array,
        lyric_attention_mask: mx.array,
        context_latents: mx.array,
        timbre_hidden_states: Optional[mx.array] = None,
    ) -> Tuple[mx.array, mx.array, mx.array]:
        """Encode all conditions into a single packed sequence.

        Args:
            text_hidden_states: [B, S_text, text_hidden_dim]
            text_attention_mask: [B, S_text]
            lyric_hidden_states: [B, S_lyric, text_hidden_dim]
            lyric_attention_mask: [B, S_lyric]
            context_latents: [B, T, 128] pre-built (src_latent + chunk_mask)
            timbre_hidden_states: optional [B, T_ref, timbre_hidden_dim]

        Returns:
            (encoder_hidden_states, encoder_attention_mask, context_latents)
        """
        encoder_hidden_states, encoder_attention_mask = self.encoder(
            text_hidden_states=text_hidden_states,
            text_attention_mask=text_attention_mask,
            lyric_hidden_states=lyric_hidden_states,
            lyric_attention_mask=lyric_attention_mask,
            timbre_hidden_states=timbre_hidden_states,
        )
        return encoder_hidden_states, encoder_attention_mask, context_latents

    def denoise_step(
        self,
        xt: mx.array,
        t: float,
        encoder_hidden_states: mx.array,
        context_latents: mx.array,
        cache: Optional[CrossAttentionCache] = None,
        debug_capture: Optional[IntermediateCapture] = None,
    ) -> mx.array:
        """Single denoising step: run the DiT to predict velocity.

        Args:
            xt: Noisy latent [B, T, acoustic_dim]
            t: Current timestep (scalar float)
            encoder_hidden_states: [B, S_enc, D] packed conditions
            context_latents: [B, T, 128]
            cache: Optional cross-attention KV cache
            debug_capture: Optional intermediate tensor capture

        Returns:
            Predicted velocity [B, T, acoustic_dim]
        """
        B = xt.shape[0]
        timestep = mx.full((B,), t)
        timestep_r = timestep  # turbo: r = t

        return self.decoder(
            hidden_states=xt,
            timestep=timestep,
            timestep_r=timestep_r,
            encoder_hidden_states=encoder_hidden_states,
            context_latents=context_latents,
            cache=cache,
            debug_capture=debug_capture,
        )
