#!/usr/bin/env python3
"""Inference-only Wisp architecture for the packaged MTP reference runtime.

This file intentionally contains no checkpoint, training, loss, data, or
release-repository dependencies.  Ship it beside ``wisp_mtp_reference.py`` in
the Hugging Face package.  Its parameter tree is identical to the training
model's inference tree.
"""

from __future__ import annotations

from dataclasses import dataclass

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


@dataclass(frozen=True)
class ModelArgs:
    vocab_size: int
    dim: int
    n_layers: int
    n_heads: int
    n_kv_heads: int
    ffn_hidden: int
    max_seq_len: int
    rope_theta: float
    norm_eps: float
    tie_embeddings: bool
    mtp_layers: int
    mtp_depth: int
    ce_chunk: int = 0

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


def causal_mask(length: int, dtype=mx.float32) -> mx.array:
    """Return Wisp's additive square causal mask."""

    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):
        batch, length, _ = x.shape
        query = self.wq(x).reshape(
            batch,
            length,
            self.n_heads,
            self.head_dim,
        )
        key = self.wk(x).reshape(
            batch,
            length,
            self.n_kv_heads,
            self.head_dim,
        )
        value = self.wv(x).reshape(
            batch,
            length,
            self.n_kv_heads,
            self.head_dim,
        )
        query = query.transpose(0, 2, 1, 3)
        key = key.transpose(0, 2, 1, 3)
        value = value.transpose(0, 2, 1, 3)

        offset = 0 if cache is None else cache[0].shape[2]
        query = self.rope(query, offset=offset)
        key = self.rope(key, offset=offset)
        if cache is not None:
            key = mx.concatenate([cache[0], key], axis=2)
            value = mx.concatenate([cache[1], value], axis=2)
        new_cache = (key, value)

        output = mx.fast.scaled_dot_product_attention(
            query,
            key,
            value,
            scale=self.scale,
            mask=mask,
        )
        output = output.transpose(0, 2, 1, 3).reshape(
            batch,
            length,
            -1,
        )
        return self.wo(output), 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, value):
        return self.w2(nn.silu(self.w1(value)) * self.w3(value))


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, value, mask=None, cache=None):
        attention, new_cache = self.attn(
            self.attn_norm(value),
            mask,
            cache,
        )
        value = value + attention
        value = value + self.ffn(self.ffn_norm(value))
        return value, new_cache


class MTPModule(nn.Module):
    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_embeddings, mask=None, caches=None):
        value = mx.concatenate(
            [
                self.h_norm(hidden),
                self.e_norm(token_embeddings),
            ],
            axis=-1,
        )
        value = self.proj(value)
        new_caches = []
        for index, block in enumerate(self.blocks):
            cache = None if caches is None else caches[index]
            value, new_cache = block(value, mask, cache)
            new_caches.append(new_cache)
        return value, 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):
        normalized = self.norm(hidden)
        if self.args.tie_embeddings:
            return self.tok_emb.as_linear(normalized)
        return self.lm_head(normalized)

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

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