"""Modulated (DiT-style) transformer pieces used by the flow models in this family.

Pixal3D's four flow checkpoints (~20GB — the bulk of its download) are pure transformers
with no sparse convolution at all, built from these blocks. LATO.2 has flow models too
(`vertex_structured_flow`, `topo_flow`), so they live in the shared core rather than in
either port.

Details that are silent when wrong, taken from upstream rather than assumed:

* In `ModulatedTransformerCrossBlock`, `norm1` and `norm3` are NON-affine but `norm2` IS
  affine (`elementwise_affine=True`). Same shapes either way.
* `MultiHeadRMSNorm` is written upstream as `F.normalize(x, dim=-1) * gamma * sqrt(dim)`.
  `F.normalize` is L2 (x/‖x‖), and multiplying by √d turns it into RMS norm —
  x/‖x‖·√d == x/rms(x). Implemented directly as RMS to avoid the double indirection.
* RoPE phases come from two places depending on the model. Pixal3D's dense `ss_flow`
  ships them precomputed as a `rope_phases` tensor; the sparse SLAT flows do not, because
  their positions are the input's own coordinates — so `rope_phases_from_coords` derives
  them. That derivation is verified to reproduce the shipped tensor to 9.6e-7.
* qk RMS norm is applied BEFORE RoPE. They do not commute.
"""

from __future__ import annotations

import math
from typing import Optional, Tuple

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


class MultiHeadRMSNorm(nn.Module):
    """Per-head RMS norm over head_dim, with a [heads, dim] gain."""

    def __init__(self, dim: int, heads: int):
        super().__init__()
        self.gamma = mx.ones((heads, dim))
        self.eps = 1e-12

    def __call__(self, x: mx.array) -> mx.array:
        # x: [..., heads, dim]
        dt = x.dtype
        f = x.astype(mx.float32)
        rms = mx.sqrt(mx.mean(f * f, axis=-1, keepdims=True) + self.eps)
        return ((f / rms) * self.gamma).astype(dt)


def apply_rope(q: mx.array, k: mx.array, phases: mx.array) -> Tuple[mx.array, mx.array]:
    """Rotate q/k by precomputed `phases`.

    `phases` is [..., head_dim/2] — one entry per rotary pair — and is normally COMPLEX.
    Upstream builds it with `torch.polar`, i.e. unit-magnitude complex numbers, and
    applies it as a complex multiply against `view_as_complex(q)`. So the rotation's
    cos/sin are the phase's REAL and IMAGINARY parts; taking `cos(phases)` of a complex
    phase would be silently, completely wrong.

    A real-valued `phases` is also accepted and treated as an angle, since that is what
    a from-scratch RoPE would hand you.

    Pairs are (even, odd) along the last axis — `view_as_complex` reshapes to [..., N, 2]
    and treats consecutive elements as (real, imag), which is interleaved, not half-split.

    Broadcasting: q/k are [B, N, H, D]; phases of [N, D/2] is reshaped to [1, N, 1, D/2]
    so it lines up on the token axis rather than accidentally on the head axis.
    """
    if phases.dtype in (mx.complex64,):
        cos, sin = mx.real(phases), mx.imag(phases)
    else:
        cos, sin = mx.cos(phases), mx.sin(phases)

    if cos.ndim == 2 and q.ndim == 4:
        cos = cos[None, :, None, :]
        sin = sin[None, :, None, :]

    def rot(t: mx.array) -> mx.array:
        dt = t.dtype
        f = t.astype(mx.float32)
        a, b = f[..., 0::2], f[..., 1::2]
        ra = a * cos - b * sin
        rb = a * sin + b * cos
        out = mx.stack([ra, rb], axis=-1).reshape(f.shape)
        return out.astype(dt)

    return rot(q), rot(k)


class TimestepEmbedder(nn.Module):
    """Sinusoidal timestep embedding -> 2-layer MLP, as in DiT."""

    def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
        super().__init__()
        self.frequency_embedding_size = frequency_embedding_size
        self.mlp_0 = nn.Linear(frequency_embedding_size, hidden_size)
        self.mlp_2 = nn.Linear(hidden_size, hidden_size)

    @staticmethod
    def timestep_embedding(t: mx.array, dim: int, max_period: int = 10000) -> mx.array:
        half = dim // 2
        freqs = mx.exp(
            -math.log(max_period) * mx.arange(half, dtype=mx.float32) / half
        )
        args = t.astype(mx.float32)[:, None] * freqs[None]
        emb = mx.concatenate([mx.cos(args), mx.sin(args)], axis=-1)
        if dim % 2:
            emb = mx.concatenate([emb, mx.zeros((emb.shape[0], 1))], axis=-1)
        return emb

    def __call__(self, t: mx.array) -> mx.array:
        e = self.timestep_embedding(t, self.frequency_embedding_size)
        return self.mlp_2(nn.silu(self.mlp_0(e)))


class _LN(nn.Module):
    """LayerNorm computed in fp32, optional affine — upstream's LayerNorm32."""

    def __init__(self, dim: int, affine: bool, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.affine = affine
        if affine:
            self.weight = mx.ones((dim,))
            self.bias = mx.zeros((dim,))

    def __call__(self, x: mx.array) -> mx.array:
        dt = x.dtype
        f = x.astype(mx.float32)
        f = (f - mx.mean(f, -1, keepdims=True)) * mx.rsqrt(
            mx.var(f, -1, keepdims=True) + self.eps
        )
        if self.affine:
            f = f * self.weight + self.bias
        return f.astype(dt)


class DiTAttention(nn.Module):
    """Self or cross attention with optional RoPE and per-head q/k RMS norm."""

    def __init__(
        self,
        channels: int,
        num_heads: int,
        ctx_channels: Optional[int] = None,
        attn_type: str = "self",
        qkv_bias: bool = True,
        use_rope: bool = False,
        qk_rms_norm: bool = False,
    ):
        super().__init__()
        self.channels, self.num_heads = channels, num_heads
        self.head_dim = channels // num_heads
        self.scale = self.head_dim**-0.5
        self._type = attn_type
        self.use_rope = use_rope
        self.qk_rms_norm = qk_rms_norm
        ctx = ctx_channels if ctx_channels is not None else channels
        if attn_type == "self":
            self.to_qkv = nn.Linear(channels, channels * 3, bias=qkv_bias)
        else:
            self.to_q = nn.Linear(channels, channels, bias=qkv_bias)
            self.to_kv = nn.Linear(ctx, channels * 2, bias=qkv_bias)
        if qk_rms_norm:
            self.q_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
            self.k_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
        self.to_out = nn.Linear(channels, channels)

    def __call__(
        self,
        x: mx.array,
        context: Optional[mx.array] = None,
        phases: Optional[mx.array] = None,
    ) -> mx.array:
        b, n, _ = x.shape
        h, d = self.num_heads, self.head_dim
        if self._type == "self":
            qkv = self.to_qkv(x).reshape(b, n, 3, h, d)
            q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
        else:
            if context is None:
                raise ValueError("cross attention needs a context")
            m = context.shape[1]
            q = self.to_q(x).reshape(b, n, h, d)
            kv = self.to_kv(context).reshape(b, m, 2, h, d)
            k, v = kv[:, :, 0], kv[:, :, 1]
        # ORDER MATTERS: upstream applies qk RMS norm FIRST, then RoPE. These do not
        # commute — RMS norm applies a per-component gain (gamma) while RoPE rotates
        # within each (even, odd) pair, so gain-then-rotate mixes different components
        # than rotate-then-gain. Reversed, a single block still correlates 0.9998 with
        # the reference, which compounds to 0.84 over 30 blocks: silent, and only
        # findable by diffing against upstream.
        if self.qk_rms_norm:
            q, k = self.q_rms_norm(q), self.k_rms_norm(k)
        if self.use_rope and phases is not None:
            q, k = apply_rope(q, k, phases)
        q = q.transpose(0, 2, 1, 3)
        k = k.transpose(0, 2, 1, 3)
        v = v.transpose(0, 2, 1, 3)
        o = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale)
        return self.to_out(o.transpose(0, 2, 1, 3).reshape(b, n, self.channels))


class ProjectAttention(nn.Module):
    """Cross-attention to global image features, plus a projection of view-aligned ones.

    Pixal3D's `image_attn_mode="proj"`. `context` carries two things: `global`
    ([B,M,ctx]) attended over normally, and `proj` ([B,N,proj_in]) which is view-aligned
    — already one token per spatial position — so it is projected and ADDED rather than
    attended. That addition is where the pixel-alignment actually enters the model.
    """

    def __init__(self, cross_attn_block: "DiTAttention", channels: int, proj_in: int):
        super().__init__()
        self.cross_attn_block = cross_attn_block
        self.proj_linear = nn.Linear(proj_in, channels)

    def __call__(self, x: mx.array, context) -> mx.array:
        if isinstance(context, dict):
            g, pr = context["global"], context["proj"]
        else:
            g, pr = context
        return self.proj_linear(pr) + self.cross_attn_block(x, g)


class ModulatedTransformerCrossBlock(nn.Module):
    """AdaLN-modulated self-attn -> cross-attn -> FFN.

    Modulation is six chunks (shift/scale/gate for MSA and MLP). With `share_mod` the
    block holds its own `modulation` parameter that is ADDED to the incoming `mod`;
    otherwise it derives them via its own `adaLN_modulation`.

    Only the self-attention and MLP are modulated — the cross-attention path is not,
    which is why `norm2` is the affine one.
    """

    def __init__(
        self,
        channels: int,
        ctx_channels: int,
        num_heads: int,
        mlp_ratio: float = 4.0,
        share_mod: bool = False,
        use_rope: bool = False,
        qk_rms_norm: bool = False,
        qk_rms_norm_cross: bool = False,
        image_attn_mode: str = "cross",
        proj_in_channels: Optional[int] = None,
    ):
        super().__init__()
        self.share_mod = share_mod
        self.image_attn_mode = image_attn_mode
        self.norm1 = _LN(channels, affine=False, eps=1e-6)
        self.norm2 = _LN(channels, affine=True, eps=1e-6)
        self.norm3 = _LN(channels, affine=False, eps=1e-6)
        self.self_attn = DiTAttention(
            channels, num_heads, use_rope=use_rope, qk_rms_norm=qk_rms_norm
        )
        _cross = DiTAttention(
            channels,
            num_heads,
            ctx_channels=ctx_channels,
            attn_type="cross",
            qk_rms_norm=qk_rms_norm_cross,
        )
        if image_attn_mode == "proj":
            self.cross_attn = ProjectAttention(
                _cross, channels, proj_in_channels or ctx_channels
            )
        else:
            self.cross_attn = _cross
        hidden = int(channels * mlp_ratio)
        self.mlp_0 = nn.Linear(channels, hidden)
        self.mlp_2 = nn.Linear(hidden, channels)
        if share_mod:
            self.modulation = mx.zeros((6 * channels,))
        else:
            self.adaLN_modulation_1 = nn.Linear(channels, 6 * channels)

    def __call__(
        self,
        x: mx.array,
        mod: mx.array,
        context: mx.array,
        phases: Optional[mx.array] = None,
    ) -> mx.array:
        if self.share_mod:
            m = (self.modulation + mod).astype(mod.dtype)
        else:
            m = self.adaLN_modulation_1(nn.silu(mod))
        c = m.shape[-1] // 6
        sh_msa, sc_msa, g_msa, sh_mlp, sc_mlp, g_mlp = (
            m[..., i * c : (i + 1) * c] for i in range(6)
        )

        h = self.norm1(x) * (1 + sc_msa[:, None]) + sh_msa[:, None]
        x = x + self.self_attn(h, phases=phases) * g_msa[:, None]

        x = x + self.cross_attn(self.norm2(x), context)

        h = self.norm3(x) * (1 + sc_mlp[:, None]) + sh_mlp[:, None]
        h = self.mlp_2(nn.gelu_approx(self.mlp_0(h)))
        return x + h * g_mlp[:, None]


def rope_phases_from_coords(
    coords: mx.array,
    head_dim: int,
    dim: int = 3,
    rope_freq: Tuple[float, float] = (1.0, 10000.0),
) -> mx.array:
    """Derive complex RoPE phases from integer positions — [N, head_dim//2].

    Pixal3D's `ss_flow` ships these precomputed as a `rope_phases` tensor, but the
    sparse SLAT flows do NOT: their positions are the SparseTensor's own coordinates,
    which vary per input, so phases must be built at call time from `coords[:, 1:]`.

    Reproduces `RotaryPositionEmbedder`:
        freq_dim = head_dim // 2 // dim
        freqs    = rope_freq[0] / rope_freq[1] ** (arange(freq_dim) / freq_dim)
        phases   = polar(1, outer(indices.flatten(), freqs))  -> [N*dim, freq_dim]
        reshape to [N, dim*freq_dim]

    dim*freq_dim usually falls SHORT of head_dim//2 (3*21 = 63 vs 64 at head_dim 128),
    and upstream right-pads the remainder with 1+0j — a no-op rotation. Dropping that
    pad would silently shift every later pair by one slot.
    """
    if head_dim % 2:
        raise ValueError(f"head_dim must be even, got {head_dim}")
    freq_dim = head_dim // 2 // dim
    f = mx.arange(freq_dim, dtype=mx.float32) / freq_dim
    freqs = rope_freq[0] / (rope_freq[1] ** f)

    idx = coords.astype(mx.float32)  # [N, dim]
    n = idx.shape[0]
    ang = idx.reshape(-1, 1) * freqs.reshape(1, -1)  # [N*dim, freq_dim]
    ang = ang.reshape(n, dim * freq_dim)

    phases = mx.cos(ang).astype(mx.float32) + 1j * mx.sin(ang).astype(mx.float32)
    want = head_dim // 2
    if phases.shape[-1] < want:
        pad = mx.ones((n, want - phases.shape[-1]), dtype=mx.complex64)
        phases = mx.concatenate([phases, pad], axis=-1)
    return phases.astype(mx.complex64)


# --------------------------------------------------------------- sparse variants
# The SLAT flows are the same DiT, but their tokens are a SparseTensor's voxels rather
# than a dense grid. Attention runs WITHIN each batch item (never across), and RoPE
# positions come from the tensor's own coordinates.


class SparseDiTAttention(nn.Module):
    """DiT attention over a SparseTensor: per-batch-item, optional RoPE + qk RMS norm."""

    def __init__(
        self,
        channels: int,
        num_heads: int,
        ctx_channels: Optional[int] = None,
        attn_type: str = "self",
        qkv_bias: bool = True,
        use_rope: bool = False,
        qk_rms_norm: bool = False,
    ):
        super().__init__()
        self.channels, self.num_heads = channels, num_heads
        self.head_dim = channels // num_heads
        self.scale = self.head_dim**-0.5
        self._type = attn_type
        self.use_rope = use_rope
        self.qk_rms_norm = qk_rms_norm
        ctx = ctx_channels if ctx_channels is not None else channels
        if attn_type == "self":
            self.to_qkv = nn.Linear(channels, channels * 3, bias=qkv_bias)
        else:
            self.to_q = nn.Linear(channels, channels, bias=qkv_bias)
            self.to_kv = nn.Linear(ctx, channels * 2, bias=qkv_bias)
        if qk_rms_norm:
            self.q_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
            self.k_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
        self.to_out = nn.Linear(channels, channels)

    def __call__(self, x, context=None, phases: Optional[mx.array] = None):
        n = x.feats.shape[0]
        h, d = self.num_heads, self.head_dim
        if self._type == "self":
            qkv = self.to_qkv(x.feats).reshape(n, 3, h, d)
            q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2]
            lq = lkv = x.layout
        else:
            cf = context.feats if hasattr(context, "feats") else context
            m = cf.shape[0]
            q = self.to_q(x.feats).reshape(n, h, d)
            kv = self.to_kv(cf).reshape(m, 2, h, d)
            k, v = kv[:, 0], kv[:, 1]
            lq = x.layout
            lkv = context.layout if hasattr(context, "layout") else [slice(0, m)]

        # Same ordering rule as the dense path: RMS norm first, then rotate.
        if self.qk_rms_norm:
            q, k = self.q_rms_norm(q), self.k_rms_norm(k)
        if self.use_rope and phases is not None and self._type == "self":
            q4, k4 = apply_rope(q[None], k[None], phases)
            q, k = q4[0], k4[0]

        outs = []
        for sq, skv in zip(lq, lkv):
            qi = q[sq].transpose(1, 0, 2)[None]
            ki = k[skv].transpose(1, 0, 2)[None]
            vi = v[skv].transpose(1, 0, 2)[None]
            o = mx.fast.scaled_dot_product_attention(qi, ki, vi, scale=self.scale)
            outs.append(o[0].transpose(1, 0, 2))
        o = outs[0] if len(outs) == 1 else mx.concatenate(outs, axis=0)
        return x.replace(self.to_out(o.reshape(n, self.channels)))


class SparseProjectAttention(nn.Module):
    """Sparse `image_attn_mode="proj"`: cross-attend `global`, add projected `proj`."""

    def __init__(self, cross_attn_block, channels: int, proj_in: int):
        super().__init__()
        self.cross_attn_block = cross_attn_block
        self.proj_linear = nn.Linear(proj_in, channels)

    def __call__(self, x, context):
        if isinstance(context, dict):
            g, pr = context["global"], context["proj"]
        else:
            g, pr = context
        prf = pr.feats if hasattr(pr, "feats") else pr
        return x.replace(self.proj_linear(prf) + self.cross_attn_block(x, g).feats)


class ModulatedSparseTransformerCrossBlock(nn.Module):
    """Sparse counterpart of ModulatedTransformerCrossBlock.

    Modulation is per BATCH ITEM but features are a flat [N, C] stack, so each row must
    pick up its own item's shift/scale/gate — hence the per-layout scatter rather than a
    simple broadcast. Broadcasting a [B, C] modulation against [N, C] would either fail
    or, worse, silently apply item 0's modulation to everything when B == 1.
    """

    def __init__(
        self,
        channels: int,
        ctx_channels: int,
        num_heads: int,
        mlp_ratio: float = 4.0,
        share_mod: bool = False,
        use_rope: bool = False,
        qk_rms_norm: bool = False,
        qk_rms_norm_cross: bool = False,
        image_attn_mode: str = "cross",
        proj_in_channels: Optional[int] = None,
    ):
        super().__init__()
        self.share_mod = share_mod
        self.norm1 = _LN(channels, affine=False, eps=1e-6)
        self.norm2 = _LN(channels, affine=True, eps=1e-6)
        self.norm3 = _LN(channels, affine=False, eps=1e-6)
        self.self_attn = SparseDiTAttention(
            channels, num_heads, use_rope=use_rope, qk_rms_norm=qk_rms_norm
        )
        _cross = SparseDiTAttention(
            channels,
            num_heads,
            ctx_channels=ctx_channels,
            attn_type="cross",
            qk_rms_norm=qk_rms_norm_cross,
        )
        self.cross_attn = (
            SparseProjectAttention(_cross, channels, proj_in_channels or ctx_channels)
            if image_attn_mode == "proj"
            else _cross
        )
        hidden = int(channels * mlp_ratio)
        self.mlp_0 = nn.Linear(channels, hidden)
        self.mlp_2 = nn.Linear(hidden, channels)
        if share_mod:
            self.modulation = mx.zeros((6 * channels,))
        else:
            self.adaLN_modulation_1 = nn.Linear(channels, 6 * channels)

    @staticmethod
    def _expand(chunk: mx.array, layout) -> mx.array:
        """[B, C] modulation -> [N, C], each row taking its own batch item's values."""
        parts = [mx.broadcast_to(chunk[i][None], (sl.stop - sl.start, chunk.shape[-1]))
                 for i, sl in enumerate(layout)]
        return parts[0] if len(parts) == 1 else mx.concatenate(parts, axis=0)

    def __call__(self, x, mod: mx.array, context, phases: Optional[mx.array] = None):
        m = (self.modulation + mod) if self.share_mod else self.adaLN_modulation_1(
            nn.silu(mod)
        )
        c = m.shape[-1] // 6
        lay = x.layout
        sh_msa, sc_msa, g_msa, sh_mlp, sc_mlp, g_mlp = (
            self._expand(m[..., i * c : (i + 1) * c], lay) for i in range(6)
        )

        h = x.replace(self.norm1(x.feats) * (1 + sc_msa) + sh_msa)
        x = x.replace(x.feats + self.self_attn(h, phases=phases).feats * g_msa)

        x = x.replace(x.feats + self.cross_attn(x.replace(self.norm2(x.feats)), context).feats)

        h = self.norm3(x.feats) * (1 + sc_mlp) + sh_mlp
        h = self.mlp_2(nn.gelu_approx(self.mlp_0(h)))
        return x.replace(x.feats + h * g_mlp)