"""Penny (DeepSeek-V4.1-Flash port) for Hugging Face transformers, via trust_remote_code.

    from transformers import AutoModelForCausalLM, AutoTokenizer
    tok = AutoTokenizer.from_pretrained("LocalWisdom/Penny-Alpha")
    model = AutoModelForCausalLM.from_pretrained("LocalWisdom/Penny-Alpha", trust_remote_code=True, dtype="float32")
    ids = tok(tok.eos_token + "The history of the bicycle begins", add_special_tokens=False, return_tensors="pt")
    out = model.generate(**ids, max_new_tokens=200, do_sample=True, temperature=0.8, top_k=50, use_cache=False)
    print(tok.decode(out[0]))

Inference-only port of the training model (no KV cache: every step re-runs the whole sequence, which is
fine at this size). Weights use the official DeepSeek-V4.1 tensor names, so `model.safetensors` loads as is.
Penny is a base model (no chat template). Prompts must start with EOS (marks a fresh document); the tokenizer does not add it, so prepend tok.eos_token.
"""

import math
from dataclasses import dataclass

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
from transformers import AutoTokenizer, GenerationMixin, PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import CausalLMOutput


# ---------------------------------------------------------------------------------------------
# config


class DeepseekV41Config(PretrainedConfig):
    """Wraps the nested config.json: top-level HF fields + `text_config` (the architecture) + `penny`."""

    model_type = "deepseek_v41"

    def __init__(self, text_config=None, penny=None, **kwargs):
        super().__init__(**kwargs)
        # kept as `arch`, not `text_config`: transformers treats a `text_config` attribute as a sub-config object
        self.arch = dict(text_config or {})
        self.penny = dict(penny or {})

    def to_dict(self):
        d = super().to_dict()
        d["text_config"] = d.pop("arch", self.arch)
        return d


@dataclass
class Args:
    vocab_size: int
    dim: int
    moe_inter_dim: int
    n_layers: int
    n_heads: int
    n_routed_experts: int
    n_shared_experts: int
    n_activated_experts: int
    score_func: str
    norm_topk_prob: bool
    route_scale: float
    swiglu_limit: float
    q_lora_rank: int
    head_dim: int
    rope_head_dim: int
    norm_eps: float
    o_groups: int
    o_lora_rank: int
    window_size: int
    compress_ratios: tuple
    kv_source_layers: tuple
    index_source_layers: tuple
    compress_rope_theta: float
    original_seq_len: int
    rope_theta: float
    rope_factor: float
    beta_fast: int
    beta_slow: int
    index_n_heads: int
    index_head_dim: int
    index_topk: int
    candidate_source_layer: int
    candidate_topk_blocks: int
    candidate_block_size: int
    hc_mult: int
    hc_sinkhorn_iters: int
    hc_eps: float
    engram_layer_ids: tuple
    engram_num_embeddings: tuple
    engram_max_ngram_size: int
    engram_vocab_size: int
    engram_n_heads: int
    engram_head_dim: int
    engram_pad_id: int
    engram_compressed_vocab_size: int
    max_seq_len: int
    attn_chunk: int = 256

    @classmethod
    def from_config(cls, cfg: DeepseekV41Config) -> "Args":
        t = cfg.arch
        rs = t.get("rope_scaling") or {}
        n = t["num_hidden_layers"]
        return cls(
            vocab_size=t["vocab_size"], dim=t["hidden_size"], moe_inter_dim=t["moe_intermediate_size"],
            n_layers=n, n_heads=t["num_attention_heads"], n_routed_experts=t["n_routed_experts"],
            n_shared_experts=t["n_shared_experts"], n_activated_experts=t["num_experts_per_tok"],
            score_func=t["scoring_func"], norm_topk_prob=t["norm_topk_prob"],
            route_scale=t["routed_scaling_factor"], swiglu_limit=t["swiglu_limit"],
            q_lora_rank=t["q_lora_rank"], head_dim=t["head_dim"], rope_head_dim=t["qk_rope_head_dim"],
            norm_eps=t["rms_norm_eps"], o_groups=t["o_groups"], o_lora_rank=t["o_lora_rank"],
            window_size=t["sliding_window"], compress_ratios=tuple(t["compress_ratios"][:n]),
            kv_source_layers=tuple(t["kv_source_layer_ids"]), index_source_layers=tuple(t["index_source_layer_ids"]),
            compress_rope_theta=t["compress_rope_theta"],
            original_seq_len=rs.get("original_max_position_embeddings", 0), rope_theta=t["rope_theta"],
            rope_factor=rs.get("factor", 1), beta_fast=rs.get("beta_fast", 32), beta_slow=rs.get("beta_slow", 1),
            index_n_heads=t["index_n_heads"], index_head_dim=t["index_head_dim"], index_topk=t["index_topk"],
            candidate_source_layer=t["candidate_source_layer_id"], candidate_topk_blocks=t["candidate_topk_blocks"],
            candidate_block_size=t["candidate_block_size"], hc_mult=t["hc_mult"],
            hc_sinkhorn_iters=t["hc_sinkhorn_iters"], hc_eps=t["hc_eps"],
            engram_layer_ids=tuple(t["engram_layer_ids"]), engram_num_embeddings=tuple(t["engram_num_embeddings"]),
            engram_max_ngram_size=t["engram_max_ngram_size"], engram_vocab_size=t["engram_vocab_size"],
            engram_n_heads=t["engram_n_heads"], engram_head_dim=t["engram_head_dim"],
            engram_pad_id=t["engram_pad_token_id"], engram_compressed_vocab_size=t["engram_compressed_vocab_size"],
            max_seq_len=cfg.penny.get("trained_seq_len", 8192),
        )


# ---------------------------------------------------------------------------------------------
# Engram n-gram hashing (hash ids match the DeepSeek reference exactly)


def _is_prime(n: int) -> bool:
    if n < 2:
        return False
    i = 2
    while i * i <= n:
        if n % i == 0:
            return False
        i += 1
    return True


def _next_prime(start: int, seen: set) -> int:
    c = start + 1
    while not _is_prime(c) or c in seen:
        c += 1
    return c


def _compressed_token_map(tokenizer):
    """Collapse tokens that normalize alike (" The", "the", "THE") onto one compressed id."""
    from tokenizers import Regex, normalizers

    sentinel = "\ue000"  # private-use char, as in the reference
    normalizer = normalizers.Sequence([
        normalizers.NFKC(), normalizers.NFD(), normalizers.StripAccents(), normalizers.Lowercase(),
        normalizers.Replace(Regex(r"[ \t\r\n]+"), " "), normalizers.Replace(Regex(r"^ $"), sentinel),
        normalizers.Strip(), normalizers.Replace(sentinel, " "),
    ])
    backend = tokenizer.backend_tokenizer
    key_to_new, lookup = {}, [0] * len(tokenizer)
    for token_id in range(len(tokenizer)):
        text = backend.decode([token_id], skip_special_tokens=False)
        if "\ufffd" in text:
            key = backend.id_to_token(token_id)
        else:
            normalized = normalizer.normalize_str(text)
            key = normalized if normalized else text
        new_id = key_to_new.get(key)
        if new_id is None:
            new_id = len(key_to_new)
            key_to_new[key] = new_id
        lookup[token_id] = new_id
    return lookup, len(key_to_new)


class NgramHash:
    """input_ids [B, L] -> hash ids [B, L, n_engram_layers, (max_ngram_size-1) * n_heads]. Built lazily
    (needs the tokenizer) and not an nn.Module, so it never touches the weights."""

    def __init__(self, a: Args, tokenizer):
        layer_ids = tuple(a.engram_layer_ids)
        self.max_n = a.engram_max_ngram_size
        primes, seen = [], set()
        for _ in layer_ids:
            per_ngram = []
            for _ in range(self.max_n - 1):
                sizes, current = [], a.engram_vocab_size - 1
                for _ in range(a.engram_n_heads):
                    current = _next_prime(current, seen)
                    seen.add(current)
                    sizes.append(current)
                per_ngram.append(sizes)
            primes.append(per_ngram)
        token_map, vocab_size = _compressed_token_map(tokenizer)
        assert vocab_size == a.engram_compressed_vocab_size, (vocab_size, a.engram_compressed_vocab_size)
        self.pad_id = token_map[a.engram_pad_id]
        flat = [[p for per_ngram in layer for p in per_ngram] for layer in primes]
        offsets = [np.cumsum([0, *sizes[:-1]]) for sizes in flat]
        max_long = np.iinfo(np.int64).max
        bound = max(1, (max_long // vocab_size) // 2)
        mult = []
        for layer_id in layer_ids:
            g = np.random.default_rng(10007 * layer_id)
            mult.append(torch.tensor(g.integers(low=0, high=bound, size=(self.max_n,), dtype=np.int64) * 2 + 1))
        self.primes = torch.tensor(primes)
        self.offsets = torch.tensor(np.array(offsets))
        self.multipliers = torch.stack(mult)
        self.token_map = torch.tensor(token_map)

    @torch.no_grad()
    def __call__(self, input_ids: torch.Tensor) -> torch.Tensor:
        dev = input_ids.device
        primes, offsets, multipliers, token_map = (
            self.primes.to(dev), self.offsets.to(dev), self.multipliers.to(dev), self.token_map.to(dev))
        compressed = token_map[input_ids]
        tokens = [compressed]
        for shift in range(1, self.max_n):
            shifted = torch.full_like(compressed, self.pad_id)
            shifted[:, shift:] = compressed[:, :-shift]
            tokens.append(shifted)
        tokens = torch.stack(tokens, dim=-1)
        products = tokens.unsqueeze(2) * multipliers
        rolling, hashes = products[..., 0], []
        for i in range(1, self.max_n):
            rolling = torch.bitwise_xor(rolling, products[..., i])
            hashes.append(rolling.unsqueeze(-1) % primes[:, i - 1])
        return torch.cat(hashes, dim=-1) + offsets


# ---------------------------------------------------------------------------------------------
# primitives


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        dtype = x.dtype
        x = x.float()
        x = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps)
        return (self.weight.float() * x).to(dtype)


def precompute_freqs_cis(dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow, device):
    freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
    if original_seq_len > 0:
        def corrected_dim(rotations):
            return dim * math.log(original_seq_len / (rotations * 2 * math.pi)) / (2 * math.log(base))

        low = max(math.floor(corrected_dim(beta_fast)), 0)
        high = min(math.ceil(corrected_dim(beta_slow)), dim - 1)
        ramp = ((torch.arange(dim // 2, dtype=torch.float32) - low) / max(high - low, 1e-3)).clamp(0, 1)
        smooth = 1 - ramp
        freqs = freqs / factor * (1 - smooth) + freqs * smooth
    freqs = torch.outer(torch.arange(seqlen), freqs)
    return torch.polar(torch.ones_like(freqs), freqs).to(device)


def rope(x, freqs_cis, rd, inverse=False):
    """Rotate the last `rd` dims of x (adjacent pairs as complex numbers). x: [b,s,d] or [b,s,h,d]."""
    nope, pe = x[..., :-rd], x[..., -rd:]
    z = torch.view_as_complex(pe.float().unflatten(-1, (-1, 2)).contiguous())
    f = freqs_cis.conj() if inverse else freqs_cis
    f = f.view(1, z.size(1), *([1] * (z.ndim - 3)), z.size(-1))
    pe = torch.view_as_real(z * f).flatten(-2).to(x.dtype)
    return torch.cat([nope, pe], dim=-1)


def idxs_to_mask(idxs, n):
    """[b,m,k] int indices (-1 = empty) -> [b,m,n] bool mask."""
    mask = torch.zeros(*idxs.shape[:2], n + 1, dtype=torch.bool, device=idxs.device)
    mask.scatter_(-1, torch.where(idxs < 0, n, idxs).long(), True)
    return mask[..., :n]


def chunked_sink_attention(q, kv, attn_sink, window, scale, comp_kv=None, comp_idxs=None, chunk=256):
    """Sliding-window attention over raw KV plus top-k compressed KV, with a learned sink logit.
    q [b,L,h,d]; kv [b,L,d]; comp_kv [b,T,d]; comp_idxs [b,L,k] into comp_kv (-1 = empty)."""
    L = q.size(1)
    outs = []
    for s0 in range(0, L, chunk):
        s1 = min(s0 + chunk, L)
        k0 = max(0, s0 - window + 1)
        qpos = torch.arange(s0, s1, device=q.device).unsqueeze(1)
        kpos = torch.arange(k0, s1, device=q.device)
        wmask = (kpos <= qpos) & (kpos > qpos - window)
        qc, kw = q[:, s0:s1], kv[:, k0:s1]
        s = torch.einsum("bmhd,bnd->bhmn", qc, kw).float() * scale
        s = s.masked_fill(~wmask, -torch.inf)
        vals = kw
        if comp_kv is not None:
            cmask = idxs_to_mask(comp_idxs[:, s0:s1], comp_kv.size(1))
            sc = torch.einsum("bmhd,bnd->bhmn", qc, comp_kv).float() * scale
            s = torch.cat([s, sc.masked_fill(~cmask.unsqueeze(1), -torch.inf)], dim=-1)
            vals = torch.cat([kw, comp_kv], dim=1)
        sink = attn_sink.float().view(1, -1, 1, 1).expand(*s.shape[:3], 1)
        lse = torch.logsumexp(torch.cat([s, sink], dim=-1), dim=-1)
        p = torch.exp(s - lse.unsqueeze(-1))
        outs.append(torch.einsum("bhmn,bnd->bmhd", p.to(vals.dtype), vals))
    return torch.cat(outs, dim=1)


def hc_split_sinkhorn(mixes, hc_scale, hc_base, hc, iters, eps):
    pre = torch.sigmoid(mixes[..., :hc] * hc_scale[0] + hc_base[:hc]) + eps
    post = 2 * torch.sigmoid(mixes[..., hc : 2 * hc] * hc_scale[1] + hc_base[hc : 2 * hc])
    comb = mixes[..., 2 * hc :] * hc_scale[2] + hc_base[2 * hc :]
    comb = comb.unflatten(-1, (hc, hc)).softmax(-1) + eps
    comb = comb / (comb.sum(-2, keepdim=True) + eps)
    for _ in range(iters - 1):
        comb = comb / (comb.sum(-1, keepdim=True) + eps)
        comb = comb / (comb.sum(-2, keepdim=True) + eps)
    return pre, post, comb


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


class Shared:
    """Per-forward hand-off between source layers and the layers that reuse their results."""

    def __init__(self):
        self.compress_kv = None
        self.index_k = None
        self.topk_idxs = None
        self.candidates = None


class Compressor(nn.Module):
    """Pools `ratio` consecutive tokens into one KV latent with a learned softmax gate (pre-RoPE)."""

    def __init__(self, a: Args, layer_id: int):
        super().__init__()
        self.ratio = a.compress_ratios[layer_id]
        self.norm = RMSNorm(a.head_dim, a.norm_eps)
        self.wkv = nn.Linear(a.dim, a.head_dim, bias=False)
        self.wgate = nn.Linear(a.dim, a.head_dim, bias=False) if self.ratio > 1 else None

    def forward(self, x):
        if self.ratio == 1:
            return self.norm(self.wkv(x))
        cutoff = x.size(1) - x.size(1) % self.ratio  # trailing partial group is not visible yet
        xf = x[:, :cutoff].float()  # the reference pools in fp32
        kv = F.linear(xf, self.wkv.weight.float()).unflatten(1, (-1, self.ratio))
        score = F.linear(xf, self.wgate.weight.float()).unflatten(1, (-1, self.ratio))
        kv = (kv * score.softmax(dim=2)).sum(dim=2)
        return self.norm(kv.to(x.dtype))


def select_candidate_blocks(logits, compress_lens, topk_blocks, block_size):
    width = logits.size(-1)
    scores = F.pad(logits, (0, -width % block_size), value=-torch.inf)
    scores = scores.unflatten(-1, (-1, block_size)).amax(dim=-1)
    num_blocks = scores.size(-1)
    last = (compress_lens - 1) // block_size
    scores = scores.masked_fill(torch.arange(num_blocks, device=logits.device) == last, torch.inf)
    top = scores.topk(min(topk_blocks, num_blocks), dim=-1)
    keep = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, top.indices, top.values > -torch.inf)
    return keep.repeat_interleave(block_size, dim=-1)[..., :width]


class Indexer(nn.Module):
    """Scores compressed positions and keeps the `index_topk` best per query."""

    def __init__(self, a: Args, layer_id: int):
        super().__init__()
        self.owns_k = layer_id in a.kv_source_layers
        self.ratio = a.compress_ratios[layer_id]
        self.is_candidate_source = layer_id == a.candidate_source_layer
        self.uses_candidates = 0 <= a.candidate_source_layer < layer_id
        self.candidate_topk_blocks = a.candidate_topk_blocks
        self.candidate_block_size = a.candidate_block_size
        self.n_heads = a.index_n_heads
        self.index_head_dim = a.index_head_dim
        self.rd = a.rope_head_dim
        self.index_topk = a.index_topk
        self.softmax_scale = self.index_head_dim**-0.5
        self.wq_b = nn.Linear(a.q_lora_rank, self.n_heads * self.index_head_dim, bias=False)
        self.weights_proj = nn.Linear(a.dim, self.n_heads, bias=False)
        if self.owns_k:
            self.wk = nn.Linear(a.head_dim, self.index_head_dim, bias=False)
            self.k_norm = RMSNorm(self.index_head_dim, a.norm_eps)

    @torch.no_grad()
    def forward(self, x, qr, latent, freqs_cis, shared: Shared, chunk: int):
        """Returns selected compressed indices [b,s,k] (-1 = empty)."""
        seqlen, ratio = x.size(1), self.ratio
        if self.owns_k:
            k = self.k_norm(self.wk(latent))
            shared.index_k = rope(k, freqs_cis[: seqlen - seqlen % ratio : ratio], self.rd)
        q = rope(self.wq_b(qr).unflatten(-1, (self.n_heads, self.index_head_dim)), freqs_cis[:seqlen], self.rd)
        index_k = shared.index_k
        weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads**-0.5)
        T = index_k.size(1)
        topk = min(self.index_topk, T)
        idx_chunks, cand_chunks = [], []
        for s0 in range(0, seqlen, chunk):
            s1 = min(s0 + chunk, seqlen)
            score = torch.einsum("bshd,btd->bsht", q[:, s0:s1], index_k)
            score = (score.relu() * weights[:, s0:s1].unsqueeze(-1)).sum(dim=2).float()
            lens = (torch.arange(s0 + 1, s1 + 1, device=x.device) // ratio).unsqueeze(-1)
            score = score.masked_fill(torch.arange(T, device=x.device) >= lens, -torch.inf)
            if self.is_candidate_source:
                cand_chunks.append(select_candidate_blocks(
                    score, lens, self.candidate_topk_blocks, self.candidate_block_size))
            elif self.uses_candidates and shared.candidates is not None:
                score = score.masked_fill(~shared.candidates[:, s0:s1], -torch.inf)
            idxs = score.topk(topk, dim=-1, sorted=False).indices
            idx_chunks.append(torch.where(idxs < lens, idxs, -1))
        if self.is_candidate_source:
            shared.candidates = torch.cat(cand_chunks, dim=1)
        return torch.cat(idx_chunks, dim=1)


class Attention(nn.Module):
    """Sliding window over raw KV, plus (compress_ratio > 0) top-k compressed KV further back.
    Low-rank Q; grouped low-rank output projection; shared K=V latent (MQA)."""

    def __init__(self, layer_id: int, a: Args):
        super().__init__()
        self.n_heads, self.head_dim, self.rd = a.n_heads, a.head_dim, a.rope_head_dim
        self.n_groups, self.o_lora_rank, self.window = a.o_groups, a.o_lora_rank, a.window_size
        self.ratio = a.compress_ratios[layer_id]
        self.softmax_scale = self.head_dim**-0.5
        self.attn_chunk = a.attn_chunk
        self.attn_sink = nn.Parameter(torch.zeros(self.n_heads))
        self.wq_a = nn.Linear(a.dim, a.q_lora_rank, bias=False)
        self.q_norm = RMSNorm(a.q_lora_rank, a.norm_eps)
        self.wq_b = nn.Linear(a.q_lora_rank, self.n_heads * self.head_dim, bias=False)
        self.wkv = nn.Linear(a.dim, self.head_dim, bias=False)
        self.kv_norm = RMSNorm(self.head_dim, a.norm_eps)
        self.wo_a = nn.Linear(self.n_heads * self.head_dim // self.n_groups, self.n_groups * self.o_lora_rank, bias=False)
        self.wo_b = nn.Linear(self.n_groups * self.o_lora_rank, a.dim, bias=False)
        self.is_kv_source = layer_id in a.kv_source_layers
        self.is_index_source = layer_id in a.index_source_layers
        self.compressor = Compressor(a, layer_id) if self.is_kv_source else None
        self.indexer = Indexer(a, layer_id) if self.is_index_source else None
        # RoPE table: built lazily on the right device (never part of the checkpoint)
        if self.ratio:
            self._rope = (a.original_seq_len, a.compress_rope_theta)
        else:
            self._rope = (0, a.rope_theta)  # pure sliding window: no YaRN, base theta
        self._rope_args = (a.rope_factor, a.beta_fast, a.beta_slow, a.max_seq_len)
        self._fc = None

    def freqs(self, seqlen, device):
        if self._fc is None or self._fc.size(0) < seqlen or self._fc.device != device:
            factor, bf, bs, max_len = self._rope_args
            self._fc = precompute_freqs_cis(self.rd, max(seqlen, max_len), self._rope[0], self._rope[1],
                                            factor, bf, bs, device)
        return self._fc[:seqlen]

    def forward(self, x, shared: Shared):
        bsz, seqlen, _ = x.shape
        fc = self.freqs(seqlen, x.device)
        qr = self.q_norm(self.wq_a(x))
        q = rope(self.wq_b(qr).unflatten(-1, (self.n_heads, self.head_dim)), fc, self.rd)
        kv = rope(self.kv_norm(self.wkv(x)), fc, self.rd)

        comp_kv, cidx = None, None
        if self.ratio:
            latent = self.compressor(x) if self.is_kv_source else None
            if self.is_index_source:
                shared.topk_idxs = self.indexer(x, qr, latent, fc, shared, chunk=self.attn_chunk)
            cidx = shared.topk_idxs
            if latent is not None:  # group j sits at the position of its first token, j * ratio
                shared.compress_kv = rope(latent, fc[: seqlen - seqlen % self.ratio : self.ratio], self.rd)
            comp_kv = shared.compress_kv

        o = chunked_sink_attention(q, kv, self.attn_sink, self.window, self.softmax_scale, comp_kv, cidx,
                                   chunk=self.attn_chunk)
        o = rope(o, fc, self.rd, inverse=True)  # remove the query rotation from the output
        o = o.reshape(bsz, seqlen, self.n_groups, -1)
        wo_a = self.wo_a.weight.view(self.n_groups, self.o_lora_rank, -1)
        o = torch.einsum("bsgd,grd->bsgr", o, wo_a.to(o.dtype))
        return self.wo_b(o.flatten(2))


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


class Gate(nn.Module):
    """sqrtsoftplus scores; a bias picks experts, weights come from the unbiased scores."""

    def __init__(self, a: Args):
        super().__init__()
        self.topk, self.score_func = a.n_activated_experts, a.score_func
        self.norm_topk_prob, self.route_scale = a.norm_topk_prob, a.route_scale
        self.weight = nn.Parameter(torch.empty(a.n_routed_experts, a.dim))
        self.register_buffer("bias", torch.zeros(a.n_routed_experts))

    def forward(self, x):
        scores = F.linear(x.float(), self.weight.float())
        if self.score_func == "softmax":
            scores = scores.softmax(dim=-1)
        elif self.score_func == "sigmoid":
            scores = scores.sigmoid()
        else:
            scores = F.softplus(scores).sqrt()
        indices = (scores + self.bias.float()).topk(self.topk, dim=-1)[1]
        weights = scores.gather(1, indices)
        if self.norm_topk_prob and self.topk > 1:
            weights = weights / (weights.sum(dim=-1, keepdim=True) + 1e-20)
        return weights * self.route_scale, indices


class Expert(nn.Module):
    def __init__(self, dim, inter_dim, swiglu_limit):
        super().__init__()
        self.w1 = nn.Linear(dim, inter_dim, bias=False)
        self.w2 = nn.Linear(inter_dim, dim, bias=False)
        self.w3 = nn.Linear(dim, inter_dim, bias=False)
        self.swiglu_limit = swiglu_limit

    def forward(self, x, weights=None):
        dtype = x.dtype
        gate, up = self.w1(x).float(), self.w3(x).float()
        if self.swiglu_limit > 0:
            up = up.clamp(-self.swiglu_limit, self.swiglu_limit)
            gate = gate.clamp(max=self.swiglu_limit)
        h = F.silu(gate) * up
        if weights is not None:
            h = weights * h
        return self.w2(h.to(dtype))


class MoE(nn.Module):
    def __init__(self, a: Args):
        super().__init__()
        self.dim = a.dim
        self.gate = Gate(a)
        self.experts = nn.ModuleList([Expert(a.dim, a.moe_inter_dim, a.swiglu_limit) for _ in range(a.n_routed_experts)])
        assert a.n_shared_experts == 1
        self.shared_experts = Expert(a.dim, a.moe_inter_dim, a.swiglu_limit)

    def forward(self, x):
        shape = x.shape
        x = x.reshape(-1, self.dim)
        weights, indices = self.gate(x)
        y = torch.zeros_like(x, dtype=torch.float32)
        counts = torch.bincount(indices.flatten(), minlength=len(self.experts)).tolist()
        for i, expert in enumerate(self.experts):
            if counts[i] == 0:
                continue
            idx, top = torch.where(indices == i)
            y = y.index_add(0, idx, expert(x[idx], weights[idx, top, None]).float())
        y = y + self.shared_experts(x).float()
        return y.to(x.dtype).view(shape)


# ---------------------------------------------------------------------------------------------
# Engram, blocks, model


class Engram(nn.Module):
    """Gated n-gram lookup written into each hc copy of the residual stream."""

    def __init__(self, a: Args, layer_id: int):
        super().__init__()
        self.layer_hash_index = tuple(a.engram_layer_ids).index(layer_id)
        self.dim, self.hc_mult, self.eps = a.dim, a.hc_mult, a.norm_eps
        self.embed = nn.Embedding(a.engram_num_embeddings[self.layer_hash_index], a.engram_head_dim)
        n_hash_cols = (a.engram_max_ngram_size - 1) * a.engram_n_heads
        self.wkv = nn.Linear(n_hash_cols * a.engram_head_dim, a.dim * (a.hc_mult + 1), bias=False)
        self.q_weight = nn.Parameter(torch.ones(a.hc_mult, a.dim))
        self.k_weight = nn.Parameter(torch.ones(a.hc_mult, a.dim))

    def forward(self, x, hash_ids):
        kv = self.wkv(self.embed(hash_ids).flatten(-2))
        key, value = kv.split([self.hc_mult * self.dim, self.dim], dim=-1)
        key = key.float().unflatten(-1, (self.hc_mult, self.dim))
        h = x.float()
        rstd = torch.rsqrt(h.square().mean(-1) + self.eps) * torch.rsqrt(key.square().mean(-1) + self.eps)
        dot = (h * (self.q_weight.float() * self.k_weight.float()) * key).sum(-1) * rstd * self.dim**-0.5
        gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(1e-6).sqrt(), dot))
        return (h + gate.unsqueeze(-1) * value.float().unsqueeze(-2)).to(x.dtype)


class Block(nn.Module):
    """Hyper-Connections block: the residual stream is carried as hc_mult copies."""

    def __init__(self, layer_id: int, a: Args):
        super().__init__()
        self.norm_eps = a.norm_eps
        self.attn = Attention(layer_id, a)
        self.ffn = MoE(a)
        self.engram = Engram(a, layer_id) if layer_id in a.engram_layer_ids else None
        self.attn_norm = RMSNorm(a.dim, a.norm_eps)
        self.ffn_norm = RMSNorm(a.dim, a.norm_eps)
        self.hc, self.hc_iters, self.hc_eps = a.hc_mult, a.hc_sinkhorn_iters, a.hc_eps
        mix_hc = (2 + self.hc) * self.hc
        self.hc_attn_fn = nn.Parameter(torch.empty(mix_hc, self.hc * a.dim))
        self.hc_ffn_fn = nn.Parameter(torch.empty(mix_hc, self.hc * a.dim))
        self.hc_attn_base = nn.Parameter(torch.empty(mix_hc))
        self.hc_ffn_base = nn.Parameter(torch.empty(mix_hc))
        self.hc_attn_scale = nn.Parameter(torch.ones(3))
        self.hc_ffn_scale = nn.Parameter(torch.ones(3))

    def hc_mixes(self, x, fn, scale, base):
        x = x.flatten(2).float()
        mixes = F.linear(x, fn.float()) * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.norm_eps)
        return hc_split_sinkhorn(mixes, scale.float(), base.float(), self.hc, self.hc_iters, self.hc_eps)

    @staticmethod
    def hc_pre(x, pre_mix):
        return torch.sum(pre_mix.unsqueeze(-1) * x.float(), dim=2).to(x.dtype)

    @staticmethod
    def hc_post(x, residual, post, comb):
        y = post.unsqueeze(-1) * x.unsqueeze(-2) + torch.sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim=2)
        return y.to(residual.dtype)

    def forward(self, x, pre_mix, shared: Shared):
        residual = x
        attn_pre, attn_post, attn_comb = self.hc_mixes(x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base)
        h = self.attn(self.attn_norm(self.hc_pre(x, pre_mix)), shared)
        x = self.hc_post(h, residual, attn_post, attn_comb)
        residual = x
        ffn_pre, ffn_post, ffn_comb = self.hc_mixes(x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base)
        h = self.ffn(self.ffn_norm(self.hc_pre(x, attn_pre)))
        x = self.hc_post(h, residual, ffn_post, ffn_comb)
        return x, ffn_pre


class DeepseekV41PreTrainedModel(PreTrainedModel):
    config_class = DeepseekV41Config
    base_model_prefix = ""
    supports_gradient_checkpointing = False

    def _init_weights(self, module):
        pass  # every parameter comes from model.safetensors


class DeepseekV41ForCausalLM(DeepseekV41PreTrainedModel, GenerationMixin):
    """Submodule names (embed, layers, norm, head) match the checkpoint's tensor names directly."""

    def __init__(self, config: DeepseekV41Config):
        super().__init__(config)
        a = Args.from_config(config)
        self.args = a
        self.embed = nn.Embedding(a.vocab_size, a.dim)
        self.layers = nn.ModuleList([Block(i, a) for i in range(a.n_layers)])
        self.norm = RMSNorm(a.dim, a.norm_eps)
        self.head = nn.Linear(a.dim, a.vocab_size, bias=False)
        self.hc = a.hc_mult
        self._engram_hash = None
        self.post_init()

    def get_input_embeddings(self):
        return self.embed

    def set_input_embeddings(self, value):
        self.embed = value

    def get_output_embeddings(self):
        return self.head

    def set_output_embeddings(self, value):
        self.head = value

    def _engram(self):
        if self._engram_hash is None:  # needs the tokenizer shipped in the same repo / folder
            src = getattr(self.config, "_name_or_path", None) or self.config.name_or_path
            self._engram_hash = NgramHash(self.args, AutoTokenizer.from_pretrained(src))
        return self._engram_hash

    def forward(self, input_ids=None, attention_mask=None, labels=None, logits_to_keep=0, **kwargs):
        """No KV cache; attention_mask is ignored (use one unpadded sequence per row, i.e. batch of equal
        lengths or batch size 1). logits_to_keep=N returns logits for the last N positions only."""
        bsz, seqlen = input_ids.shape
        hashes = self._engram()(input_ids)
        h = self.embed(input_ids).unsqueeze(2).repeat(1, 1, self.hc, 1)
        pre_mix = torch.zeros(bsz, seqlen, self.hc, device=h.device)
        pre_mix[..., 0] = 1.0
        shared = Shared()
        for layer in self.layers:
            if layer.engram is not None:
                h = layer.engram(h, hashes[:, :, layer.engram.layer_hash_index])
            h, pre_mix = layer(h, pre_mix, shared)
        h = self.norm(Block.hc_pre(h, pre_mix))
        logits = self.head(h[:, -logits_to_keep:] if logits_to_keep else h)
        loss = None
        if labels is not None:
            loss = F.cross_entropy(logits[:, :-1].float().flatten(0, 1), labels[:, 1:].flatten(), ignore_index=-100)
        return CausalLMOutput(loss=loss, logits=logits)