"""
Wisp model: a Llama-style decoder in MLX with a native multi-token-prediction
module attached, in the Qwen3-Next / DeepSeek-V3 style.

Design notes
------------
* One shared MTP module, applied recursively for depth > 1. This matches what
  MTPLX expects at inference (draft depth is a runtime knob, not a parameter
  count), and it is far cheaper than DeepSeek's one-module-per-depth layout.
* The MTP module consumes (trunk hidden state at position i, embedding of the
  token at position i+1) and predicts the token at position i+2. Recursion feeds
  the module's own output back in as the hidden state.
* The LM head is shared between the trunk and the MTP module. Sharing ties both
  to one output projection, which is cheap and removes a whole set of parameters
  that could drift apart. It does not by itself force the drafter's distribution
  close to the target's: the hidden states feeding that shared head are produced
  by different computations. Whether the distributions are actually close is an
  empirical question, and acceptance rate is the measurement of it.
"""

from dataclasses import dataclass, asdict

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


@dataclass
class ModelArgs:
    vocab_size: int = 32768
    dim: int = 768
    n_layers: int = 12
    n_heads: int = 12
    n_kv_heads: int = 4
    ffn_hidden: int = 2048
    max_seq_len: int = 2048
    rope_theta: float = 100000.0
    norm_eps: float = 1e-5
    tie_embeddings: bool = True
    mtp_layers: int = 1
    mtp_depth: int = 2
    ce_chunk: int = 0          # 0 disables chunking, else rows per chunk

    @property
    def head_dim(self) -> int:
        return self.dim // self.n_heads

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_dict(cls, d: dict) -> "ModelArgs":
        known = {k: v for k, v in d.items() if k in cls.__dataclass_fields__}
        return cls(**known)


def causal_mask(length: int, dtype=mx.float32) -> mx.array:
    """Additive causal mask of shape (length, length)."""
    upper = mx.triu(mx.ones((length, length), dtype=mx.bool_), k=1)
    return mx.where(upper, mx.array(-1e9, dtype=dtype), mx.array(0.0, dtype=dtype))


class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.n_heads = args.n_heads
        self.n_kv_heads = args.n_kv_heads
        self.head_dim = args.head_dim
        self.scale = self.head_dim ** -0.5

        self.wq = nn.Linear(args.dim, args.n_heads * args.head_dim, bias=False)
        self.wk = nn.Linear(args.dim, args.n_kv_heads * args.head_dim, bias=False)
        self.wv = nn.Linear(args.dim, args.n_kv_heads * args.head_dim, bias=False)
        self.wo = nn.Linear(args.n_heads * args.head_dim, args.dim, bias=False)
        self.rope = nn.RoPE(args.head_dim, traditional=False, base=args.rope_theta)

    def __call__(self, x, mask=None, cache=None):
        b, length, _ = x.shape

        q = self.wq(x).reshape(b, length, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
        k = self.wk(x).reshape(b, length, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
        v = self.wv(x).reshape(b, length, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3)

        offset = 0 if cache is None else cache[0].shape[2]
        q = self.rope(q, offset=offset)
        k = self.rope(k, offset=offset)

        if cache is not None:
            k = mx.concatenate([cache[0], k], axis=2)
            v = mx.concatenate([cache[1], v], axis=2)
        new_cache = (k, v)

        # mx.fast.scaled_dot_product_attention natively supports grouped query
        # attention and explicitly documents that k and v should not be
        # pre-tiled to match q's head count. The previous mx.repeat here
        # materialized k and v at the full head count before every attention
        # call, in every layer, every micro-step: with n_heads 12 and
        # n_kv_heads 4 that is a 3x larger tensor than the fused kernel needs,
        # pure wasted memory bandwidth. Verified bit-identical output against
        # the tiled path before removing it (scripts/test_gqa_attention.py).
        out = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale, mask=mask)
        out = out.transpose(0, 2, 1, 3).reshape(b, length, -1)
        return self.wo(out), new_cache


class FeedForward(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.w1 = nn.Linear(args.dim, args.ffn_hidden, bias=False)
        self.w3 = nn.Linear(args.dim, args.ffn_hidden, bias=False)
        self.w2 = nn.Linear(args.ffn_hidden, args.dim, bias=False)

    def __call__(self, x):
        return self.w2(nn.silu(self.w1(x)) * self.w3(x))


class Block(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.attn_norm = nn.RMSNorm(args.dim, eps=args.norm_eps)
        self.attn = Attention(args)
        self.ffn_norm = nn.RMSNorm(args.dim, eps=args.norm_eps)
        self.ffn = FeedForward(args)

    def __call__(self, x, mask=None, cache=None):
        attn_out, new_cache = self.attn(self.attn_norm(x), mask, cache)
        x = x + attn_out
        x = x + self.ffn(self.ffn_norm(x))
        return x, new_cache


class MTPModule(nn.Module):
    """Predicts one token further ahead than whatever produced its input hidden state."""

    def __init__(self, args: ModelArgs):
        super().__init__()
        self.h_norm = nn.RMSNorm(args.dim, eps=args.norm_eps)
        self.e_norm = nn.RMSNorm(args.dim, eps=args.norm_eps)
        self.proj = nn.Linear(2 * args.dim, args.dim, bias=False)
        self.blocks = [Block(args) for _ in range(args.mtp_layers)]

    def __call__(self, hidden, token_emb, mask=None, caches=None):
        x = mx.concatenate([self.h_norm(hidden), self.e_norm(token_emb)], axis=-1)
        x = self.proj(x)
        new_caches = []
        for i, block in enumerate(self.blocks):
            cache = None if caches is None else caches[i]
            x, nc = block(x, mask, cache)
            new_caches.append(nc)
        return x, new_caches


class Wisp(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.args = args
        self.tok_emb = nn.Embedding(args.vocab_size, args.dim)
        self.blocks = [Block(args) for _ in range(args.n_layers)]
        self.norm = nn.RMSNorm(args.dim, eps=args.norm_eps)
        if not args.tie_embeddings:
            self.lm_head = nn.Linear(args.dim, args.vocab_size, bias=False)
        self.mtp = MTPModule(args)

    def head(self, hidden):
        h = self.norm(hidden)
        if self.args.tie_embeddings:
            return self.tok_emb.as_linear(h)
        return self.lm_head(h)

    def trunk(self, tokens, mask=None, caches=None):
        h = self.tok_emb(tokens)
        new_caches = []
        for i, block in enumerate(self.blocks):
            cache = None if caches is None else caches[i]
            h, nc = block(h, mask, cache)
            new_caches.append(nc)
        return h, new_caches

    def __call__(self, tokens, mask=None, caches=None):
        h, new_caches = self.trunk(tokens, mask, caches)
        return self.head(h), h, new_caches

    def cross_entropy(self, hidden, targets):
        """
        Cross entropy over flattened positions, optionally without ever holding
        the whole (N, vocab) logits tensor.

        The logits are the largest tensor in the step by a wide margin. At
        micro_batch 16 and seq_len 2048 one is 1.07GB in bfloat16, and a step
        materialises `1 + mtp_depth` of them, each of which must stay live for its
        own backward. That is why throughput barely responds to batch size: the
        step is moving bytes, not doing arithmetic.

        With `ce_chunk` set, each chunk goes through `mx.checkpoint`, so its logits
        are recomputed during the backward instead of being kept. The parameters
        are passed as explicit arguments rather than captured, because a closure
        capture would be treated as a constant and would silently drop the
        gradients for the norm and the output projection.
        """
        h = hidden.reshape(-1, self.args.dim)
        t = targets.reshape(-1)
        n = h.shape[0]
        chunk = self.args.ce_chunk

        if not chunk or chunk >= n:
            return nn.losses.cross_entropy(self.head(h), t, reduction="mean")

        w_norm = self.norm.weight
        w_out = self.tok_emb.weight if self.args.tie_embeddings else self.lm_head.weight
        eps = self.args.norm_eps

        def piece(h_, t_, wn, wo):
            hh = mx.fast.rms_norm(h_, wn, eps)
            return nn.losses.cross_entropy(hh @ wo.T, t_, reduction="sum")

        ckpt = mx.checkpoint(piece)
        total = ckpt(h[:chunk], t[:chunk], w_norm, w_out)
        for s in range(chunk, n, chunk):
            total = total + ckpt(h[s:s + chunk], t[s:s + chunk], w_norm, w_out)
        return total / n

    def loss(self, batch, mtp_weight: float = 0.3):
        """
        batch: (B, T) int32 where T = seq_len + 1 + mtp_depth.

        Returns (total, main_loss, [mtp_loss_per_depth]).

        Index bookkeeping: position i of the input sees batch[:, i] and the trunk
        predicts batch[:, i+1]. MTP step k consumes the depth-(k-1) hidden state
        plus the embedding of batch[:, i+k] and predicts batch[:, i+k+1].
        """
        depth = self.args.mtp_depth
        seq_len = batch.shape[1] - 1 - depth

        inputs = batch[:, :seq_len]
        mask = causal_mask(seq_len, inputs.dtype if inputs.dtype != mx.int32 else mx.float32)
        mask = mask.astype(self.norm.weight.dtype)

        hidden, _ = self.trunk(inputs, mask)
        main = self.cross_entropy(hidden, batch[:, 1:seq_len + 1])

        mtp_losses = []
        cur = hidden
        for k in range(1, depth + 1):
            emb = self.tok_emb(batch[:, k:seq_len + k])
            cur, _ = self.mtp(cur, emb, mask)
            mtp_losses.append(
                self.cross_entropy(cur, batch[:, k + 1:seq_len + k + 1])
            )

        total = main
        if depth > 0:
            total = main + mtp_weight * sum(mtp_losses) / depth
        return total, main, mtp_losses

    def n_params(self, trunk_only: bool = False) -> int:
        from mlx.utils import tree_flatten

        def count(tree):
            return sum(v.size for _, v in tree_flatten(tree) if isinstance(v, mx.array))

        if trunk_only:
            return count(self.tok_emb.parameters()) + count(
                [b.parameters() for b in self.blocks]
            ) + count(self.norm.parameters())
        return count(self.parameters())