"""MLX implementation of the DiffusionGemma text model (encoder + block-diffusion decoder).

Shape conventions: activations are [B, L, hidden]; attention is [B, heads, L, head_dim]; the KV cache keeps
[B, heads, S, head_dim] per layer. Weight names follow the HF checkpoint (``model.decoder.*``), which is also
where the encoder's weights live (in the checkpoint the encoder shares them). A linear weight is either a float
array ``X.weight`` or NVFP4-packed ``X.weight`` (uint32, 8 nibbles per word) + ``X.scales`` (uint8, group 16).
Batch size is 1: the port serves one request at a time, so no padding/batch masks are needed.
"""

from __future__ import annotations

import math
from typing import Any

import mlx.core as mx

from .config import DiffusionGemmaConfig, TextConfig

NVFP4 = dict(group_size=16, bits=4, mode="nvfp4")


# --------------------------------------------------------------------------- weights


def _packed(sd: dict[str, mx.array], name: str) -> tuple[mx.array, mx.array | None, mx.array | None]:
    """(weight, scales, global_scale) for a checkpoint linear name ending in ``.weight``."""
    stem = name[: -len(".weight")] if name.endswith(".weight") else name
    return sd[name], sd.get(stem + ".scales"), sd.get(stem + ".global_scale")


class Linear:
    """``x @ W.T`` for a float or NVFP4-packed weight."""

    def __init__(self, weight: mx.array, scales: mx.array | None = None, global_scale: mx.array | None = None):
        self.weight = weight
        self.scales = scales
        self.global_scale = global_scale

    def __call__(self, x: mx.array) -> mx.array:
        if self.scales is None:
            return x @ self.weight.T
        if self.global_scale is not None:
            # mlx 0.32.4 has no global_scale on quantized_matmul (only gather_qmm); no plain linear in
            # this checkpoint is packed, so this is a guard, not a path.
            raise TypeError("packed Linear with a global scale is not supported")
        return mx.quantized_matmul(x, self.weight, scales=self.scales, transpose=True, **NVFP4)


class RMSNorm:
    """DiffusionGemma RMSNorm: ``x / sqrt(mean(x^2) + eps) * w``, computed in float32."""

    def __init__(self, weight: mx.array | None, eps: float):
        self.weight = weight  # None -> no scale (v_norm, router norm, self-conditioning post norm)
        self.eps = eps

    def __call__(self, x: mx.array) -> mx.array:
        dtype = x.dtype
        x = x.astype(mx.float32)
        x = x * mx.power(mx.mean(x * x, axis=-1, keepdims=True) + self.eps, -0.5)
        if self.weight is not None:
            x = x * self.weight.astype(mx.float32)
        return x.astype(dtype)


def gelu_tanh(x: mx.array) -> mx.array:
    return 0.5 * x * (1.0 + mx.tanh(0.7978845608028654 * (x + 0.044715 * x * x * x)))


# --------------------------------------------------------------------------- attention


def rope_tables(text: TextConfig, layer_type: str, positions: mx.array) -> tuple[mx.array, mx.array]:
    """cos/sin of shape [L, head_dim/2]; both head halves rotate with the same angles (reference layout)."""
    params = text.rope_parameters[layer_type]
    head_dim = text.global_head_dim if layer_type == "full_attention" else text.head_dim
    base = params["rope_theta"]
    factor = params.get("partial_rotary_factor", 1.0)
    if params.get("rope_type", "default") == "proportional":
        # transformers' proportional RoPE: rotate the first dims, zero frequencies pad the rest.
        angles = int(factor * head_dim // 2)
        exponents = mx.arange(0, 2 * angles, 2, dtype=mx.float32) / head_dim
        inv_freq = mx.concatenate(
            [1.0 / (base**exponents), mx.zeros(head_dim // 2 - angles, dtype=mx.float32)]
        )
    else:
        dim = int(factor * head_dim)
        inv_freq = 1.0 / (base ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim))
    freqs = positions.astype(mx.float32)[:, None] * inv_freq[None, :]
    return mx.cos(freqs), mx.sin(freqs)


def apply_rope(x: mx.array, cos: mx.array, sin: mx.array) -> mx.array:
    """x [B, L, H, D], cos/sin [L, D/2]."""
    half = x.shape[-1] // 2
    cos, sin = cos[None, :, None, :], sin[None, :, None, :]
    x1, x2 = x[..., :half], x[..., half:]
    return mx.concatenate([x1 * cos - x2 * sin, x1 * sin + x2 * cos], axis=-1)


class KVCache:
    """Per-layer prefix cache. Sliding layers keep only ``sliding_window - 1`` tokens, like the reference."""

    def __init__(self, text: TextConfig):
        n = text.num_hidden_layers
        self.text = text
        self.keys: list[mx.array | None] = [None] * n
        self.values: list[mx.array | None] = [None] * n
        self.seen = [0] * n  # tokens encoded so far; can exceed the kept length on sliding layers

    def append(self, i: int, k: mx.array, v: mx.array) -> tuple[mx.array, mx.array]:
        """Append [B, H, L, D] states; returns all states visible to this forward."""
        self.seen[i] += k.shape[2]
        if self.keys[i] is not None:
            k = mx.concatenate([self.keys[i], k], axis=2)
            v = mx.concatenate([self.values[i], v], axis=2)
        if self.text.is_sliding(i):
            keep = self.text.sliding_window - 1
            self.keys[i] = k[:, :, -keep:, :] if k.shape[2] > keep else k
            self.values[i] = v[:, :, -keep:, :] if v.shape[2] > keep else v
        else:
            self.keys[i], self.values[i] = k, v
        return k, v

    def prefix(self, i: int) -> tuple[mx.array | None, mx.array | None, int]:
        """Cached K/V and the absolute position of the first cached token."""
        if self.keys[i] is None:
            return None, None, 0
        return self.keys[i], self.values[i], self.seen[i] - self.keys[i].shape[2]

    def copy(self) -> KVCache:
        """Shallow copy: ``append`` always replaces entries, so a copy can be extended without touching this one."""
        other = KVCache.__new__(KVCache)
        other.text = self.text
        other.keys = list(self.keys)
        other.values = list(self.values)
        other.seen = list(self.seen)
        return other


class Attention:
    def __init__(self, text: TextConfig, layer_idx: int, sd: dict[str, mx.array], prefix: str):
        self.layer_idx = layer_idx
        self.head_dim = text.layer_head_dim(layer_idx)
        self.num_heads = text.num_attention_heads
        self.num_kv_heads = text.layer_num_kv_heads(layer_idx)
        self.scaling = 1.0  # the reference does not scale by 1/sqrt(head_dim)
        self.q_proj = Linear(*_packed(sd, f"{prefix}.q_proj.weight"))
        self.k_proj = Linear(*_packed(sd, f"{prefix}.k_proj.weight"))
        v_proj = f"{prefix}.v_proj.weight"
        # Full-attention layers have no v_proj: V is K after k_norm + rope.
        self.v_proj = Linear(*_packed(sd, v_proj)) if v_proj in sd else None
        self.o_proj = Linear(*_packed(sd, f"{prefix}.o_proj.weight"))
        eps = text.rms_norm_eps
        self.q_norm = RMSNorm(sd[f"{prefix}.q_norm.weight"], eps)
        self.k_norm = RMSNorm(sd[f"{prefix}.k_norm.weight"], eps)
        self.v_norm = RMSNorm(None, eps)  # no scale

    def __call__(self, x, cos, sin, mask, cache: KVCache, write: bool) -> mx.array:
        b, length, _ = x.shape
        hd, nh, nkv = self.head_dim, self.num_heads, self.num_kv_heads
        cos, sin = cos.astype(x.dtype), sin.astype(x.dtype)

        q = apply_rope(self.q_norm(self.q_proj(x).reshape(b, length, nh, hd)), cos, sin).transpose(0, 2, 1, 3)
        k = self.k_proj(x).reshape(b, length, nkv, hd)
        v = k if self.v_proj is None else self.v_proj(x).reshape(b, length, nkv, hd)
        k = apply_rope(self.k_norm(k), cos, sin)
        v = self.v_norm(v)
        k, v = k.transpose(0, 2, 1, 3), v.transpose(0, 2, 1, 3)

        if write:  # encoder: append to the prefix cache
            k, v = cache.append(self.layer_idx, k, v)
        else:  # diffusion decoder: cached prefix + canvas, read-only
            pk, pv, _ = cache.prefix(self.layer_idx)
            if pk is not None:
                if pk.shape[0] != b:  # batched draws share the one encoded prefix
                    pk = mx.broadcast_to(pk, (b,) + pk.shape[1:])
                    pv = mx.broadcast_to(pv, (b,) + pv.shape[1:])
                k = mx.concatenate([pk, k], axis=2)
                v = mx.concatenate([pv, v], axis=2)

        out = mx.fast.scaled_dot_product_attention(
            q, k, v, scale=self.scaling, mask=mask
        )
        return self.o_proj(out.transpose(0, 2, 1, 3).reshape(b, length, nh * hd))


# --------------------------------------------------------------------------- MoE


class Router:
    def __init__(self, text: TextConfig, sd: dict[str, mx.array], prefix: str):
        self.top_k = text.top_k_experts
        self.scalar_root_size = text.hidden_size**-0.5
        self.norm = RMSNorm(None, text.rms_norm_eps)
        self.proj = Linear(*_packed(sd, f"{prefix}.proj.weight"))
        self.scale = sd[f"{prefix}.scale"]
        self.per_expert_scale = sd[f"{prefix}.per_expert_scale"]

    def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]:
        """x [T, hidden] -> (top-k expert ids, top-k weights), both [T, top_k]."""
        h = self.norm(x) * self.scale * self.scalar_root_size
        probs = mx.softmax(self.proj(h).astype(mx.float32), axis=-1)
        idx = mx.argpartition(-probs, kth=self.top_k, axis=-1)[..., : self.top_k]
        weights = mx.take_along_axis(probs, idx, axis=-1)
        order = mx.argsort(-weights, axis=-1)  # descending, like torch.topk
        idx = mx.take_along_axis(idx, order, axis=-1)
        weights = mx.take_along_axis(weights, order, axis=-1)
        weights = weights / weights.sum(axis=-1, keepdims=True)
        return idx, weights * self.per_expert_scale[idx]


def _expert_weights(sd: dict[str, mx.array], prefix: str, num_experts: int, kinds: list[str]):
    """Stacked (weight, scales, global_scale) over experts, one entry per kind concatenated on the output axis.

    Takes the converted artifact's pre-stacked tensors as they are (``experts.gate_up_proj`` /
    ``experts.down_proj``, with ``global_scale``) and falls back to stacking the checkpoint's per-expert
    layout (``experts.{e}.gate_proj`` etc.) for float weights, as the tiny parity tests need.
    """
    base = f"{prefix}.experts"
    if f"{base}.{kinds[0]}.weight" in sd:
        return _packed(sd, f"{base}.{kinds[0]}.weight")
    if f"{base}.{kinds[0]}" in sd:  # reference layout: one Parameter, no suffix (packed or float)
        name = f"{base}.{kinds[0]}"
        if sd[name].dtype in (mx.uint32, mx.uint8):
            return sd[name], sd.get(name + ".scales"), sd.get(name + ".global_scale")
        return sd[name], None, None
    weights, scales = [], []
    for e in range(num_experts):
        parts = [_packed(sd, f"{base}.{e}.{kind}.weight") for kind in kinds]
        weights.append(mx.concatenate([p[0] for p in parts], axis=0))
        scales.append(None if parts[0][1] is None else mx.concatenate([p[1] for p in parts], axis=0))
    return mx.stack(weights), None if scales[0] is None else mx.stack(scales), None


class Experts:
    """128-expert MLP; gate and up stacked on the output axis, as in the reference's ``gate_up_proj``."""

    def __init__(self, text: TextConfig, sd: dict[str, mx.array], prefix: str):
        n = text.num_experts
        kinds = (
            ["gate_up_proj"]
            if f"{prefix}.experts.gate_up_proj.weight" in sd or f"{prefix}.experts.gate_up_proj" in sd
            else ["gate_proj", "up_proj"]
        )
        self.gate_up, self.gate_up_scales, self.gate_up_gs = _expert_weights(sd, prefix, n, kinds)
        self.down, self.down_scales, self.down_gs = _expert_weights(sd, prefix, n, ["down_proj"])

    def __call__(self, x: mx.array, idx: mx.array, weights: mx.array) -> mx.array:
        """x [T, hidden] with idx/weights [T, top_k] -> [T, hidden].

        Expert assignments are sorted by expert before the gathers: unsorted nvfp4 gather_qmm is ~9x slower
        here (measured on this Mac, 9 ms vs 1 ms per call), because the kernel cannot stream the weights.
        """
        t, k = idx.shape
        flat = idx.reshape(-1)
        order = mx.argsort(flat)
        idx_s = mx.take(flat, order)
        xs = mx.take(x, order // k, axis=0)
        gate_up = self._mm(xs, self.gate_up, self.gate_up_scales, self.gate_up_gs, idx_s)
        gate, up = mx.split(gate_up, 2, axis=-1)
        h = gelu_tanh(gate) * up
        out = self._mm(h.reshape(t * k, -1), self.down, self.down_scales, self.down_gs, idx_s)
        out = mx.take(out, mx.argsort(order), axis=0).reshape(t, k, -1)
        # router probabilities are float32; promote the 9x bigger activation tensor instead of the weights
        return (out * weights[..., None].astype(out.dtype)).sum(axis=-2)

    @staticmethod
    def _mm(x: mx.array, w: mx.array, scales: mx.array | None, gs: mx.array | None, idx: mx.array) -> mx.array:
        """x [N, in] with sorted idx [N] expert ids -> [N, out]."""
        x = mx.expand_dims(x, -2)
        if scales is None:
            out = mx.gather_mm(x, w.swapaxes(-1, -2), rhs_indices=idx, sorted_indices=True)
        else:
            out = mx.gather_qmm(
                x, w, scales, rhs_indices=idx, transpose=True, sorted_indices=True, global_scale=gs, **NVFP4
            )
        return out[:, 0, :]


# --------------------------------------------------------------------------- transformer blocks


class Layer:
    def __init__(self, text: TextConfig, layer_idx: int, sd: dict[str, mx.array], prefix: str):
        self.self_attn = Attention(text, layer_idx, sd, f"{prefix}.self_attn")
        self.gate_proj = Linear(*_packed(sd, f"{prefix}.mlp.gate_proj.weight"))
        self.up_proj = Linear(*_packed(sd, f"{prefix}.mlp.up_proj.weight"))
        self.down_proj = Linear(*_packed(sd, f"{prefix}.mlp.down_proj.weight"))
        eps = text.rms_norm_eps
        self.input_layernorm = RMSNorm(sd[f"{prefix}.input_layernorm.weight"], eps)
        self.post_attention_layernorm = RMSNorm(sd[f"{prefix}.post_attention_layernorm.weight"], eps)
        self.pre_feedforward_layernorm = RMSNorm(sd[f"{prefix}.pre_feedforward_layernorm.weight"], eps)
        self.pre_feedforward_layernorm_2 = RMSNorm(sd[f"{prefix}.pre_feedforward_layernorm_2.weight"], eps)
        self.post_feedforward_layernorm = RMSNorm(sd[f"{prefix}.post_feedforward_layernorm.weight"], eps)
        self.post_feedforward_layernorm_1 = RMSNorm(sd[f"{prefix}.post_feedforward_layernorm_1.weight"], eps)
        self.post_feedforward_layernorm_2 = RMSNorm(sd[f"{prefix}.post_feedforward_layernorm_2.weight"], eps)
        self.layer_scalar = sd[f"{prefix}.layer_scalar"]
        self.router = Router(text, sd, f"{prefix}.router")
        self.experts = Experts(text, sd, prefix)

    def __call__(self, x, cos, sin, mask, cache: KVCache, write: bool) -> mx.array:
        attn = self.self_attn(self.input_layernorm(x), cos, sin, mask, cache, write)
        residual = x + self.post_attention_layernorm(attn)

        h = self.pre_feedforward_layernorm(residual)
        dense = self.post_feedforward_layernorm_1(self.down_proj(gelu_tanh(self.gate_proj(h)) * self.up_proj(h)))

        flat = residual.reshape(-1, residual.shape[-1])
        # Routing reads the raw residual; the experts read its pre-normalised version.
        idx, weights = self.router(flat)
        moe = self.experts(self.pre_feedforward_layernorm_2(flat), idx, weights)
        moe = self.post_feedforward_layernorm_2(moe.reshape(residual.shape))

        return (residual + self.post_feedforward_layernorm(dense + moe)) * self.layer_scalar


class SelfConditioning:
    def __init__(self, text: TextConfig, sd: dict[str, mx.array], prefix: str):
        eps = text.rms_norm_eps
        self.pre_norm = RMSNorm(sd[f"{prefix}.pre_norm.weight"], eps)
        self.post_norm = RMSNorm(None, eps)
        self.gate_proj = Linear(*_packed(sd, f"{prefix}.gate_proj.weight"))
        self.up_proj = Linear(*_packed(sd, f"{prefix}.up_proj.weight"))
        self.down_proj = Linear(*_packed(sd, f"{prefix}.down_proj.weight"))

    def __call__(self, inputs_embeds: mx.array, signal: mx.array) -> mx.array:
        h = self.pre_norm(signal)
        h = self.down_proj(gelu_tanh(self.gate_proj(h)) * self.up_proj(h))
        return self.post_norm(inputs_embeds + h)


# --------------------------------------------------------------------------- model


class DiffusionGemma:
    """Text-only DiffusionGemma: an encoder that writes a prefix KV cache, and a diffusion decoder."""

    def __init__(self, config: DiffusionGemmaConfig, state_dict: dict[str, Any]):
        text = config.text_config
        self.config = config
        self.text = text
        sd = {k: (v if isinstance(v, mx.array) else mx.array(v)) for k, v in state_dict.items()}

        self.embed_tokens = sd["model.decoder.embed_tokens.weight"]
        self.lm_head = sd.get("lm_head.weight", self.embed_tokens)
        # The reference downcasts sqrt(hidden_size) to the embedding dtype (bf16 -> 53.0 for 2816).
        self.embed_scale = mx.array(math.sqrt(text.hidden_size)).astype(self.embed_tokens.dtype)

        self.layers = [Layer(text, i, sd, f"model.decoder.layers.{i}") for i in range(text.num_hidden_layers)]
        self.norm = RMSNorm(sd["model.decoder.norm.weight"], text.rms_norm_eps)
        self.self_conditioning = SelfConditioning(text, sd, "model.decoder.self_conditioning")

    # -- encoder -----------------------------------------------------------------

    def _encoder_mask(self, layer_idx: int, positions: mx.array, kv_len: int, kv_pos0: int, block_ids):
        """Causal (+ sliding window) mask, with bidirectional groups where ``block_ids`` agree."""
        text = self.text
        kv_pos = mx.arange(kv_len) + kv_pos0
        allowed = kv_pos[None, :] <= positions[:, None]
        if text.is_sliding(layer_idx):
            allowed &= kv_pos[None, :] > (positions - text.sliding_window)[:, None]
        if block_ids is not None:
            # keys already in the cache belong to no group of this chunk (-1), so they stay causal
            pad = kv_len - block_ids.shape[0]
            keys_ids = mx.concatenate([mx.full((pad,), -1, mx.int32), block_ids]) if pad > 0 else block_ids
            allowed |= (block_ids[:, None] == keys_ids[None, :]) & (block_ids[:, None] >= 0)
        # the dtype the attention kernel wants; the mask is cached per layer type, so cast it once here
        return mx.where(allowed, 0.0, -mx.inf)[None, None].astype(self.embed_tokens.dtype)

    def encode(
        self,
        input_ids: mx.array,
        positions: mx.array,
        cache: KVCache | None = None,
        block_ids: mx.array | None = None,
    ) -> tuple[mx.array, KVCache]:
        """Encode ``input_ids`` into the prefix cache; returns (last_hidden_state, cache).

        ``block_ids`` (one int per token) marks groups that attend bidirectionally inside the group -- the
        per-question blocks of the reference's ``create_masks_for_generate``; ``-1`` is plain causal text.
        As in the reference, they apply to the prefill only (they are dropped for later canvases).
        """
        text = self.text
        cache = cache if cache is not None else KVCache(text)
        if block_ids is not None and cache.seen[0] != 0 and block_ids.shape[0] != input_ids.shape[1]:
            raise ValueError("block_ids must cover exactly the tokens being encoded")
        x = self.embed_tokens[input_ids] * self.embed_scale
        ropes, masks = {}, {}
        for i, layer in enumerate(self.layers):
            layer_type = text.layer_types[i]
            if layer_type not in ropes:
                ropes[layer_type] = rope_tables(text, layer_type, positions)
                _, _, kv_pos0 = cache.prefix(i)
                kv_len = cache.seen[i] - kv_pos0 + input_ids.shape[1]
                masks[layer_type] = self._encoder_mask(i, positions, kv_len, kv_pos0, block_ids)
            x = layer(x, *ropes[layer_type], masks[layer_type] if text.is_causal else None, cache, write=True)
        return self.norm(x), cache

    # -- decoder -----------------------------------------------------------------

    def denoise(
        self,
        canvas_ids: mx.array,
        cache: KVCache,
        positions: mx.array,
        self_conditioning_logits: mx.array | None = None,
        read_positions: mx.array | None = None,
        vocab_ids: mx.array | None = None,
    ) -> mx.array:
        """One denoiser forward over the canvas; returns softcapped logits [B, canvas, vocab] in float32.

        ``read_positions``/``vocab_ids`` project only those canvas rows onto those vocabulary ids, which is
        what a constrained read needs and saves the full-vocabulary unembedding of every canvas position.
        """
        text = self.text
        x = self.embed_tokens[canvas_ids] * self.embed_scale
        if self_conditioning_logits is not None:
            probs = mx.softmax(self_conditioning_logits.astype(mx.float32), axis=-1)
            soft = (probs.astype(self.embed_tokens.dtype) @ self.embed_tokens) * self.embed_scale
        else:
            soft = mx.zeros_like(x)
        x = self.self_conditioning(x, soft)

        ropes = {}
        for i, layer in enumerate(self.layers):
            layer_type = text.layer_types[i]
            if layer_type not in ropes:
                ropes[layer_type] = rope_tables(text, layer_type, positions)
            x = layer(x, *ropes[layer_type], None, cache, write=False)
        h = self.norm(x)
        if read_positions is not None:
            h = mx.take(h, read_positions, axis=1)
        head = self.lm_head if vocab_ids is None else mx.take(self.lm_head, vocab_ids, axis=0)
        logits = (h @ head.T).astype(mx.float32)
        return mx.tanh(logits / text.final_logit_softcapping) * text.final_logit_softcapping