"""Foundry-4: sparse hybrid kernel-optimization model.

Architecture (see DESIGN notes at bottom):

- Embedding + KernelGram n-gram memory (hashed bigram..5-gram tables,
  zero-init gates — proven on Foundry-2/3 to carry ~half the val loss).
- Hybrid trunk: GatedDeltaNet linear-attention blocks everywhere except
  every ``attn_every``-th layer which is full causal attention (RoPE +
  QK-norm + SDPA).  Linear layers handle the stream; attention layers
  give exact long-range recall for pointer names / block sizes.
- Sparse MoE FFN on every block: 1 shared expert + N fine-grained
  routed experts, top-k sigmoid router (DeepSeekMoE-style), Switch-style
  load-balance aux loss.
- PonderLoop: a recurrent block (shared weights, iteration embedding)
  applied K times to the trunk output.  Per-iteration deep supervision
  (weighted CE on each iteration's logits) so every loop step is a
  usable state — the substrate for latent "try many paths" reasoning
  and later RL-driven non-lazy search.
- JudgmentHead (Jev-like): small MLP over the pooled latent state that
  emits typed verdicts — p(compiles), p(correct), expected speedup,
  halt.  In pretrain it is self-supervised on a signal we already have:
  per-iteration LM improvement.  In RL it is re-pointed at real harness
  outcomes and becomes the best-of-loop selector / idea filter.
- Tied 65k-token embedding + chunked CE scoring (never materialise
  [B,T,V] logits at once), optional MTP head.

Invariants (same discipline as Foundry-1/2 — do not break):
  1. Trunk is causal.  Attention uses the explicit -inf mask; GDN is
     causal by construction; the ponder block uses the same mask.
  2. Rewards/judgments are measured or self-supervised, never faked.
  3. save_checkpoint stage->verify->copy->verify on every write.
"""
from __future__ import annotations

import math
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint

from gdn import GatedDeltaNet
from segmm import seg_mm, wscatter


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

@dataclass
class Foundry4Config:
    vocab_size: int = 65536          # bpe2_65k_v2 + specials 0..15
    n_specials: int = 16
    d_model: int = 1280
    n_layers: int = 24
    n_heads: int = 10                # attn heads, head_dim 128
    head_dim: int = 128
    seq_len: int = 4096
    dropout: float = 0.0
    rope_theta: float = 10000.0
    qk_norm: bool = True
    attn_every: int = 4              # layer i%4==3 -> full attention
    # gated deltanet
    gdn_heads: int = 10
    gdn_head_k: int = 128
    gdn_head_v: int = 128
    gdn_chunk: int = 32          # C=64 can overflow the WY solve (2^C growth)
    # sparse MoE
    n_experts: int = 384
    expert_hidden: int = 128
    moe_topk: int = 8
    shared_hidden: int = 640         # always-on shared expert
    lb_weight: float = 0.01
    # kernelgram
    kg_buckets: int = 262144
    kg_dim: int = 80
    kg_max_n: int = 4                # n-gram orders 2..4
    kg_probes: int = 2               # salted probes per order
    # ponder loop (latent reasoning)
    ponder_iters: int = 4
    ponder_hidden: int = 2560        # dense SwiGLU inside the loop block
    ponder_iter_weight_decay: float = 1.0   # loss weight ramp exponent
    # judgment head
    judge_hidden: int = 512
    # mtp
    use_mtp: bool = True
    mtp_weight: float = 0.05
    # memory
    moe_ckpt: bool = False           # recompute MoE dispatch in backward


def beast_config() -> Foundry4Config:
    """~5.1B total / ~570M active.  The target."""
    return Foundry4Config()


def dev_config() -> Foundry4Config:
    """~700M total / ~120M active — smoke + kernel-tuning config."""
    return Foundry4Config(
        d_model=768, n_layers=12, n_heads=6, head_dim=128,
        gdn_heads=6, gdn_head_k=128, gdn_head_v=128,
        n_experts=32, expert_hidden=192, moe_topk=4,
        shared_hidden=512, seq_len=2048,
        kg_buckets=65536, kg_dim=64, kg_max_n=4, kg_probes=2,
        ponder_iters=3, ponder_hidden=1536,
        judge_hidden=256)


# ---------------------------------------------------------------------------
# KernelGram — identical semantics to model2 (proven); kept self-contained
# ---------------------------------------------------------------------------

class KernelGram(nn.Module):
    _C1 = 0x9E3779B1
    _C2 = 0x85EBCA77
    _SEEDS = (0x0, 0x9E3779B1, 0x85EBCA77, 0xC2B2AE3D)

    def __init__(self, cfg: Foundry4Config):
        super().__init__()
        self.cfg = cfg
        self.n_grams = cfg.kg_max_n - 1
        self.n_probes = cfg.kg_probes
        self.n_slots = self.n_grams * self.n_probes
        self.tables = nn.ParameterList([
            nn.Parameter(torch.randn(cfg.kg_buckets, cfg.kg_dim)
                         / math.sqrt(cfg.kg_dim))
            for _ in range(self.n_slots)])
        self.proj = nn.ParameterList([
            nn.Parameter(torch.empty(cfg.kg_dim, cfg.d_model))
            for _ in range(self.n_slots)])
        for p in self.proj:
            nn.init.normal_(p, std=0.02)
        self.gates = nn.Parameter(torch.zeros(self.n_slots))

    def _hashes(self, ids: torch.Tensor) -> list[torch.Tensor]:
        B, T = ids.shape
        out = []
        for k_ in range(self.n_grams):
            n = k_ + 2
            for p in range(self.n_probes):
                h = torch.zeros(B, T, dtype=torch.long, device=ids.device)
                if p:
                    h[:, n - 1:] = self._SEEDS[p % len(self._SEEDS)]
                for j in range(n):
                    tok = ids[:, j:T - n + 1 + j]
                    h[:, n - 1:] = (h[:, n - 1:] * self._C1
                                    + tok * self._C2)
                out.append(h % self.cfg.kg_buckets)
        return out

    def forward(self, ids, x_emb):
        for s, (h, tab, pr) in enumerate(
                zip(self._hashes(ids), self.tables, self.proj)):
            g = torch.tanh(self.gates[s])
            x_emb = x_emb + g * (F.embedding(h, tab) @ pr)
        return x_emb


# ---------------------------------------------------------------------------
# Sparse MoE — shared expert + fine-grained routed experts (sigmoid router,
# DeepSeek-style top-k with weight renormalisation; bucketed bmm dispatch)
# ---------------------------------------------------------------------------

class SparseMoE(nn.Module):
    def __init__(self, cfg: Foundry4Config):
        super().__init__()
        d, h = cfg.d_model, cfg.expert_hidden
        E = cfg.n_experts
        self.cfg = cfg
        self.E, self.k = E, cfg.moe_topk
        self.router = nn.Linear(d, E, bias=False)
        self.w_g = nn.Parameter(torch.empty(E, d, h))
        self.w_w = nn.Parameter(torch.empty(E, d, h))
        self.w_o = nn.Parameter(torch.empty(E, h, d))
        for w in (self.w_g, self.w_w, self.w_o):
            nn.init.normal_(w, std=0.02)
        # shared expert (always on)
        self.sh_g = nn.Linear(d, cfg.shared_hidden, bias=False)
        self.sh_w = nn.Linear(d, cfg.shared_hidden, bias=False)
        self.sh_o = nn.Linear(cfg.shared_hidden, d, bias=False)

    def _swiglu(self, x, wg, ww, wo):
        return (F.silu(x @ wg) * (x @ ww)) @ wo

    def _dispatch(self, xf, top_i, w):
        """Packed-segment expert compute (segmm.py Triton grouped GEMM).
        DROPLESS (causality invariant): every routed token keeps all k
        experts — a capacity mask would couple position i's output to
        future tokens' routing and break KV-cache semantics.  Packed
        layout: memory stays at Nk rows regardless of routing skew, and
        there is no host<->device sync at all."""
        N, d = xf.shape
        counts = torch.bincount(
            top_i.reshape(-1), minlength=self.E)
        order = torch.argsort(top_i.reshape(-1))          # stable sort
        flat_t = (order // self.k)                        # token ids
        flat_w = w.reshape(-1)[order]
        offs = torch.zeros(self.E + 1, dtype=torch.int64,
                           device=xf.device)
        offs[1:] = torch.cumsum(counts, 0)
        xs = xf[flat_t]                                   # [Nk,d] packed
        dt = xs.dtype
        wg = self.w_g.to(dt)                              # cast: same-dtype
        ww = self.w_w.to(dt)                              # tl.dot operands
        wo = self.w_o.to(dt)
        h1 = seg_mm(xs, wg, offs)
        h2 = seg_mm(xs, ww, offs)
        yp = seg_mm(F.silu(h1) * h2, wo, offs)            # [Nk,d]
        out = xf.new_zeros(N, d)
        out.index_add_(0, flat_t,
                       (yp * flat_w.unsqueeze(-1)).to(out.dtype))
        return out

    def forward(self, x: torch.Tensor):
        """x [B,T,d] -> out [B,T,d], lb_loss scalar."""
        B, T, d = x.shape
        xf = x.reshape(-1, d)
        N = xf.size(0)

        logits = self.router(xf)                       # [N,E]
        scores = torch.sigmoid(logits)
        top_s, top_i = scores.topk(self.k, dim=-1)     # [N,k]
        w = top_s / top_s.sum(-1, keepdim=True).clamp_min(1e-6)

        if self.cfg.moe_ckpt and self.training and xf.requires_grad:
            out = torch.utils.checkpoint.checkpoint(
                self._dispatch, xf, top_i, w,
                use_reentrant=False)
        else:
            out = self._dispatch(xf, top_i, w)

        out = out + self.sh_o(F.silu(self.sh_g(xf)) * self.sh_w(xf))
        out = out.view(B, T, d)

        # Switch-style load balance on the full softmax
        p = F.softmax(logits, -1)
        f = torch.zeros_like(p).scatter_(
            1, top_i, 1.0 / self.k).mean(0)
        lb = self.E * (f * p.mean(0)).sum()
        return out, lb


# ---------------------------------------------------------------------------
# attention mixer + block
# ---------------------------------------------------------------------------

class RMSNorm(nn.Module):
    def __init__(self, d: int, eps: float = 1e-6):
        super().__init__()
        self.w = nn.Parameter(torch.ones(d))
        self.eps = eps

    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True)
                               + self.eps) * self.w


def _rope_cos_sin(T, hd, theta, device):
    half = hd // 2
    fr = 1.0 / (theta ** (torch.arange(0, half, device=device,
                                       dtype=torch.float32) / half))
    t = torch.arange(T, device=device, dtype=torch.float32)
    ang = torch.outer(t, fr)
    return ang.cos(), ang.sin()


def _rope(x, cos, sin):
    half = x.size(-1) // 2
    x1, x2 = x[..., :half], x[..., half:]
    c = cos[None, None]
    s = sin[None, None]
    return torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], -1)


class Attention(nn.Module):
    def __init__(self, cfg: Foundry4Config):
        super().__init__()
        d = cfg.d_model
        self.h, self.hd = cfg.n_heads, cfg.head_dim
        self.wq = nn.Linear(d, self.h * self.hd, bias=False)
        self.wk = nn.Linear(d, self.h * self.hd, bias=False)
        self.wv = nn.Linear(d, self.h * self.hd, bias=False)
        self.wo = nn.Linear(self.h * self.hd, d, bias=False)
        self.qn = RMSNorm(self.hd) if cfg.qk_norm else None
        self.kn = RMSNorm(self.hd) if cfg.qk_norm else None

    def forward(self, x, cos, sin, causal):
        B, T, _ = x.shape
        q = self.wq(x).view(B, T, self.h, self.hd).transpose(1, 2)
        k = self.wk(x).view(B, T, self.h, self.hd).transpose(1, 2)
        v = self.wv(x).view(B, T, self.h, self.hd).transpose(1, 2)
        if self.qn is not None:
            q, k = self.qn(q), self.kn(k)
        q, k = _rope(q, cos, sin), _rope(k, cos, sin)
        a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.wo(a.transpose(1, 2).reshape(B, T, -1))

    def forward_cached(self, x, cos, sin, past=None):
        """Decode-time path.  past = (k, v) post-rope from earlier calls.
        Returns (out, (k_all, v_all)) — the returned tensors become the
        next call's `past`.  cos/sin must cover only the NEW positions."""
        B, T, _ = x.shape
        q = self.wq(x).view(B, T, self.h, self.hd).transpose(1, 2)
        k = self.wk(x).view(B, T, self.h, self.hd).transpose(1, 2)
        v = self.wv(x).view(B, T, self.h, self.hd).transpose(1, 2)
        if self.qn is not None:
            q, k = self.qn(q), self.kn(k)
        q, k = _rope(q, cos, sin), _rope(k, cos, sin)
        if past is not None:
            k = torch.cat([past[0], k], 2)
            v = torch.cat([past[1], v], 2)
        if past is None or T == k.size(2):
            a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        else:
            P = k.size(2) - T
            m = torch.ones(T, k.size(2), dtype=torch.bool,
                           device=x.device)
            m[:, P:] = torch.tril(torch.ones(T, T, dtype=torch.bool,
                                             device=x.device))
            a = F.scaled_dot_product_attention(q, k, v, attn_mask=m)
        return self.wo(a.transpose(1, 2).reshape(B, T, -1)), (k, v)


class Block(nn.Module):
    """mixer (GDN or attention) + sparse MoE, pre-norm."""

    def __init__(self, cfg: Foundry4Config, use_attn: bool):
        super().__init__()
        self.use_attn = use_attn
        d = cfg.d_model
        self.norm1 = RMSNorm(d)
        self.norm2 = RMSNorm(d)
        if use_attn:
            self.mixer = Attention(cfg)
        else:
            self.mixer = GatedDeltaNet(
                d, cfg.gdn_heads, cfg.gdn_head_k, cfg.gdn_head_v,
                chunk_size=cfg.gdn_chunk)
        self.ffn = SparseMoE(cfg)

    def forward(self, h, cos, sin, causal):
        if self.use_attn:
            h = h + self.mixer(self.norm1(h), cos, sin, causal)
        else:
            h = h + self.mixer(self.norm1(h))
        ff, lb = self.ffn(self.norm2(h))
        h = h + ff
        return h, lb

    def forward_cached(self, h, cos, sin, cache):
        """cache: dict — {'kv': (k,v)} for attn, {'S': state} for GDN."""
        if self.use_attn:
            o, kv = self.mixer.forward_cached(self.norm1(h), cos, sin,
                                              cache.get("kv"))
            cache["kv"] = kv
            h = h + o
        else:
            o, S, ct = self.mixer(self.norm1(h), S=cache.get("S"),
                                  use_cache=True,
                                  conv_prev=cache.get("conv"))
            cache["S"] = S
            cache["conv"] = ct
            h = h + o
        ff, _ = self.ffn(self.norm2(h))
        return h + ff


# ---------------------------------------------------------------------------
# PonderLoop — the latent reasoning core
# ---------------------------------------------------------------------------

class PonderBlock(nn.Module):
    """One shared-weight recurrent step: causal attn + dense SwiGLU."""

    def __init__(self, cfg: Foundry4Config):
        super().__init__()
        d = cfg.d_model
        self.norm1 = RMSNorm(d)
        self.norm2 = RMSNorm(d)
        self.attn = Attention(cfg)
        self.ff_g = nn.Linear(d, cfg.ponder_hidden, bias=False)
        self.ff_w = nn.Linear(d, cfg.ponder_hidden, bias=False)
        self.ff_o = nn.Linear(cfg.ponder_hidden, d, bias=False)
        self.iter_emb = nn.Parameter(torch.zeros(cfg.ponder_iters, d))

    def forward(self, h, it, cos, sin, causal):
        h = h + self.iter_emb[it]
        h = h + self.attn(self.norm1(h), cos, sin, causal)
        x = self.norm2(h)
        h = h + self.ff_o(F.silu(self.ff_g(x)) * self.ff_w(x))
        return h

    def forward_cached(self, h, it, cos, sin, past):
        o, kv = self.attn.forward_cached(
            self.norm1(h + self.iter_emb[it]), cos, sin, past)
        h = h + self.iter_emb[it] + o
        x = self.norm2(h)
        return h + self.ff_o(F.silu(self.ff_g(x)) * self.ff_w(x)), kv


class JudgmentHead(nn.Module):
    """Typed verdicts over the pooled latent state.

    Outputs: p_compile, p_correct logits, log-speedup regression,
    halt logit.  Trained self-supervised in pretrain (per-iter LM
    improvement); re-pointed at harness rewards in SFT/RL.
    """

    def __init__(self, cfg: Foundry4Config):
        super().__init__()
        d = cfg.d_model
        self.net = nn.Sequential(
            nn.Linear(d, cfg.judge_hidden), nn.SiLU(),
            nn.Linear(cfg.judge_hidden, 4))

    def forward(self, h):                      # h [B,T,d] -> [B,4]
        pooled = h.mean(1)
        return self.net(pooled)


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

class Foundry4Model(nn.Module):
    def __init__(self, cfg: Foundry4Config | None = None):
        super().__init__()
        self.cfg = cfg = cfg or Foundry4Config()
        d = cfg.d_model
        self.tok_emb = nn.Embedding(cfg.vocab_size, d)
        self.kg = KernelGram(cfg)
        self.blocks = nn.ModuleList([
            Block(cfg, use_attn=((i + 1) % cfg.attn_every == 0))
            for i in range(cfg.n_layers)])
        self.norm = RMSNorm(d)
        self.ponder = PonderBlock(cfg)
        self.judge = JudgmentHead(cfg)
        self.head = nn.Linear(d, cfg.vocab_size, bias=False)
        self.head.weight = self.tok_emb.weight          # tied
        self.mtp = (nn.Linear(d, d, bias=False)
                    if cfg.use_mtp else None)
        self.apply(self._init)

    @staticmethod
    def _init(m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, std=0.02)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, std=0.02)

    # -- trunk -------------------------------------------------------------

    def trunk(self, ids):
        x = self.kg(ids, self.tok_emb(ids))
        dev = ids.device
        T = ids.size(1)
        cos, sin = _rope_cos_sin(T, self.cfg.head_dim,
                                 self.cfg.rope_theta, dev)
        h, lb = x, x.new_zeros(())
        for blk in self.blocks:
            h, l = blk(h, cos, sin, None)
            lb = lb + l
        return self.norm(h), lb

    # -- forward: trunk -> ponder loop -> per-iter hidden states -----------

    def forward(self, ids, return_aux=False):
        cfg = self.cfg
        h, lb = self.trunk(ids)
        T = ids.size(1)
        cos, sin = _rope_cos_sin(T, cfg.head_dim, cfg.rope_theta,
                                 ids.device)
        h_iters, judge_per_iter = [h], []
        for it in range(cfg.ponder_iters):
            h = self.ponder(h, it, cos, sin, None)
            h_iters.append(h)
            if return_aux:
                judge_per_iter.append(self.judge(h))
        if return_aux:
            # training path: h_iters drive chunked/fused CE — never
            # materialise the [B,T,V] logits here (saves ~4GB + backward)
            aux = {"lb_loss": lb * cfg.lb_weight,
                   "judge": torch.stack(judge_per_iter, 1),
                   "h_iters": h_iters}
            if self.mtp is not None:
                aux["mtp_h"] = self.mtp(h)
            return None, aux
        return self.head(h)                              # final iter

    # -- cached decode ------------------------------------------------------

    def new_cache(self):
        return {"blocks": [{} for _ in self.blocks],
                "ponder": [None] * self.cfg.ponder_iters,
                "tail": None}                     # last kg_max_n-1 ids

    def forward_cached(self, ids, pos0, cache):
        """Decode/prefill for generation.  ids [B,Tn] at absolute
        positions pos0..pos0+Tn-1.  Returns logits [B,Tn,V]; cache is
        updated in place.  Matches forward() token-for-token (causality
        invariant), just incremental."""
        cfg = self.cfg
        B, Tn = ids.shape
        dev = ids.device
        # KernelGram needs the trailing n-1 ids for the n-gram window
        tail = cache["tail"]
        kg_in = ids if tail is None else torch.cat([tail, ids], 1)
        x = self.kg(kg_in, self.tok_emb(kg_in))[:, -Tn:]
        cache["tail"] = kg_in[:, -(cfg.kg_max_n - 1):]
        cos, sin = _rope_cos_sin(pos0 + Tn, cfg.head_dim,
                                 cfg.rope_theta, dev)
        cos, sin = cos[pos0:], sin[pos0:]
        h = x
        for blk, bc in zip(self.blocks, cache["blocks"]):
            h = blk.forward_cached(h, cos, sin, bc)
        h = self.norm(h)
        for it in range(cfg.ponder_iters):
            h, cache["ponder"][it] = self.ponder.forward_cached(
                h, it, cos, sin, cache["ponder"][it])
        return self.head(h)


# ---------------------------------------------------------------------------
# parameter accounting + checkpointing (same discipline as model2)
# ---------------------------------------------------------------------------

def param_report(model: Foundry4Model) -> dict:
    cfg = model.cfg
    total = sum(p.numel() for p in model.parameters())

    emb = cfg.vocab_size * cfg.d_model                       # tied
    kg = sum(p.numel() for p in model.kg.parameters())
    experts = cfg.n_layers * cfg.n_experts * 3 * cfg.d_model \
        * cfg.expert_hidden

    # active = everything except the cold (non-selected) routed-expert
    # capacity.  Embedding lookup counts once (tied head reuses it for
    # logits, so it's active anyway).
    routed = cfg.n_layers * cfg.moe_topk * 3 * cfg.d_model \
        * cfg.expert_hidden
    active = total - experts + routed
    n_attn = cfg.n_layers // cfg.attn_every
    n_gdn = cfg.n_layers - n_attn
    return {"total": total, "active": active,
            "experts": experts, "kg": kg, "emb": emb,
            "n_attn": n_attn, "n_gdn": n_gdn}


def save_checkpoint(obj: dict, out_path) -> None:
    """Stage-write-verify-copy-verify; survives flaky filesystems."""
    import io
    import shutil
    import tempfile
    from pathlib import Path

    out_path = Path(out_path)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    buf = io.BytesIO()
    torch.save(obj, buf)
    data = buf.getvalue()
    with tempfile.NamedTemporaryFile(
            dir=out_path.parent, delete=False) as f:
        f.write(data)
        tmp = Path(f.name)
    # verify the staged file loads before it replaces anything
    torch.load(tmp, map_location="cpu", weights_only=False)
    shutil.copy2(tmp, out_path)
    torch.load(out_path, map_location="cpu", weights_only=False)
    tmp.unlink()


# ---------------------------------------------------------------------------
# FLOP model for MFU accounting
# ---------------------------------------------------------------------------

def flops_per_token(cfg: Foundry4Config) -> float:
    """Forward+backward FLOPs per token (~6N_active + attention terms)."""
    d = cfg.d_model
    n_attn = cfg.n_layers // cfg.attn_every
    n_gdn = cfg.n_layers - n_attn
    # dense per-token params: mixers + shared expert + topk experts +
    # router + ponder(iters) + heads — times 6
    mixer = n_attn * 4 * d * cfg.n_heads * cfg.head_dim \
        + n_gdn * d * (4 * cfg.gdn_heads * cfg.gdn_head_k)
    ffn = cfg.n_layers * 3 * d * (cfg.shared_hidden
                                + cfg.moe_topk * cfg.expert_hidden)
    ponder = cfg.ponder_iters * (4 * d * cfg.n_heads * cfg.head_dim
                                 + 3 * d * cfg.ponder_hidden)
    heads = cfg.vocab_size * d                       # tied head proj
    n_active = mixer + ffn + ponder + heads
    # attention QK/AV term: ~4 * T * d * n_attn_heads-ish; and GDN chunk
    # work ~= O(C * dk * dv) per token — lump as seq term:
    seq_term = n_attn * 12 * cfg.seq_len * d \
        + n_gdn * 2 * cfg.gdn_chunk * cfg.gdn_heads \
        * cfg.gdn_head_v
    return 6.0 * n_active + seq_term