"""MiniV41: the assembled causal encoder-decoder model.

    tokens
      |
      v
    embed  -> [engram writes at encoder layers] -> causal ENCODER
      |
      v
    H_encoder ---> global K/V projection -----+
      |                                       |
      v                                       v
    DECODER (own sliding-window K/V  +  global K/V)
      |
      v
    logits
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field

import torch
import torch.nn.functional as F
from torch import nn

from .attention import RMSNorm, RotaryEmbedding
from .ced import Block, CEDCache, GlobalKVProjection
from .config import ModelConfig
from .engram import Engram, EngramLayout, NgramHasher, build_backend
from .execution import LayerExecutionPlan
from .experts_grouped import GroupedExpertBank
from .router import Router

__all__ = ["MiniV41", "ModelOutput"]


@dataclass
class ModelOutput:
    logits: torch.Tensor
    loss: torch.Tensor | None = None
    lm_loss: torch.Tensor | None = None
    aux_loss: torch.Tensor | None = None
    z_loss: torch.Tensor | None = None
    cache: CEDCache | None = None
    hidden_encoder: torch.Tensor | None = None
    router_stats: list[dict] = field(default_factory=list)
    engram_stats: list[dict] = field(default_factory=list)


class MiniV41(nn.Module):
    #: Router and Engram statistics for the training logs. See `set_telemetry`.
    collect_telemetry: bool = True

    def __init__(self, cfg: ModelConfig, tokenizer=None):
        super().__init__()
        self.cfg = cfg
        self.embed = nn.Embedding(cfg.vocab_size, cfg.hidden_size)

        self.rope = RotaryEmbedding(cfg.head_dim, cfg.max_sequence_length, cfg.rope_theta)

        # ---- encoder ------------------------------------------------------ #
        self.encoder = nn.ModuleList(
            [
                Block(cfg, layer_id=i, window_size=cfg.encoder_window_size, is_decoder=False)
                for i in range(cfg.encoder_layers)
            ]
        )

        # ---- engram ------------------------------------------------------- #
        self.engram_layout: EngramLayout | None = None
        self.hasher: NgramHasher | None = None
        self.engrams = nn.ModuleDict()
        if cfg.engram.enabled:
            backend = build_backend(cfg.engram, cfg.vocab_size, tokenizer)
            self.engram_backend_identity = backend.identity()
            self.engram_layout = EngramLayout.build(cfg.engram, backend.compressed_vocab_size)
            self.hasher = NgramHasher(cfg.engram, self.engram_layout, backend)
            for idx, layer_id in enumerate(self.engram_layout.layer_ids):
                self.engrams[str(layer_id)] = Engram(
                    self.engram_layout, idx, cfg.hidden_size, cfg.rms_norm_eps
                )
        else:
            self.engram_backend_identity = "disabled"

        # ---- CED bridge --------------------------------------------------- #
        if cfg.ced_enabled:
            self.enc_out_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
            if cfg.decoder_kv_source == "shared":
                self.global_kv = GlobalKVProjection(cfg)
            else:
                self.global_kv = nn.ModuleList(
                    [GlobalKVProjection(cfg) for _ in range(cfg.decoder_layers)]
                )
        else:
            self.enc_out_norm = None
            self.global_kv = None

        # ---- decoder ------------------------------------------------------ #
        self.decoder = nn.ModuleList(
            [
                Block(
                    cfg,
                    layer_id=cfg.encoder_layers + i,
                    # Without CED there is no global path, so a windowed decoder
                    # would simply be unable to see far back. Fall back to full
                    # causal so the ablation stays a fair comparison.
                    window_size=cfg.decoder_window_size if cfg.ced_enabled else 0,
                    is_decoder=True,
                )
                for i in range(cfg.decoder_layers)
            ]
        )

        # Execution plans. The decoder is always sequential: looping decoder
        # layers would change the receptive field R = W + (D-1)(W-1), which is
        # written in terms of layer count, so that is a separate change.
        # `recurrent` returns `sequential` itself when loop_count == 1, so the
        # control arm is the baseline path and not a lookalike.
        self.encoder_plan = LayerExecutionPlan.recurrent(
            cfg.encoder_layers,
            span=cfg.recurrent_span or (0, cfg.encoder_layers),
            loop_count=cfg.loop_count,
            residual_scale=cfg.recurrent_residual_scale,
            forbidden_layers=set(cfg.engram.layer_ids) if cfg.engram.enabled else set(),
        )
        self.decoder_plan = LayerExecutionPlan.sequential(cfg.decoder_layers)

        self.final_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
        if cfg.tie_word_embeddings:
            self.lm_head.weight = self.embed.weight

        # Reset by forward/prefill/decode and appended to by _collect_route.
        # Initialised here too, so run_encoder/run_decoder can be called
        # directly on a fresh model without an AttributeError.
        self._aux_losses: list[torch.Tensor] = []
        self._z_losses: list[torch.Tensor] = []

        if cfg.deterministic_init:
            self.init_weights_by_name(cfg.init_seed)
        else:
            self.apply(self._init_weights)

        # AFTER initialisation, deliberately. Zeroing before it would simply be
        # overwritten by the init that follows -- and an ablation that silently
        # does nothing is worse than no ablation, because the arm would look
        # like the control and the factorial would read as "experts do not
        # matter". Placed here so it overwrites real values.
        if cfg.moe.enabled and cfg.moe.ablate_routed_experts_to_zero:
            n = self._ablate_routed_experts()
            assert n > 0, "ablation requested but no routed expert bank found"
            # Re-init after apply() so Engram keeps its own (near-zero) scheme:
            # apply() re-drew both the table (an nn.Embedding) and wkv (an
            # nn.Linear) with the generic std, so both must be restored.
            for engram in self.engrams.values():
                nn.init.normal_(engram.table.weight, std=0.02)
                nn.init.normal_(engram.wkv.weight, std=0.02 / math.sqrt(2))

    def set_telemetry(self, enabled: bool) -> MiniV41:
        """Turn the per-forward router/Engram statistics on or off.

        On by default: `train.py` logs them and the MoE health gate reads them.
        Off is for inference only. The logits do not change -- the statistics
        are read off the forward, never fed back -- but each one is a `.item()`,
        and ~250 host syncs per decoded token made generation ~4x slower than
        the GPU work it does. `router_stats`/`engram_stats` come back empty and
        `aux_loss`/`z_loss` are zero while it is off.
        """
        self.collect_telemetry = enabled
        for mod in self.modules():
            if isinstance(mod, Router):
                mod.collect_stats = enabled
        return self

    @torch.no_grad()
    def _ablate_routed_experts(self) -> int:
        """Zero every routed expert bank, deliberately. Returns how many.

        Called AFTER initialisation, so it overwrites real values rather than
        depending on an uninitialised tensor happening to be zero -- which is
        what made the V1 accident invisible. Nothing holds the banks down
        afterwards because nothing needs to: zero is absorbing for SwiGLU.
        """
        n = 0
        for mod in self.modules():
            bank = getattr(mod, "grouped", None)
            if bank is not None:
                bank.w1.zero_()
                bank.w3.zero_()
                bank.w2.zero_()
                n += 1
                continue
            experts = getattr(mod, "experts", None)
            if experts:
                for e in experts:
                    for w in ("w1", "w2", "w3"):
                        getattr(e, w).weight.zero_()
                n += 1
        return n

    def _init_weights(self, module: nn.Module) -> None:
        std = self.cfg.initializer_range
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=std)
        elif isinstance(module, GroupedExpertBank):
            # The grouped bank holds RAW nn.Parameters, not nn.Linear submodules,
            # so the two branches above walked straight past it and every routed
            # expert stayed at `torch.empty`, which on fresh pages is zeros.
            #
            # Zero is an absorbing state for SwiGLU: with w1 = w3 = w2 = 0 the
            # output is zero, so dL/dw2 is zero (the activation is zero) and
            # dL/dw1, dL/dw3 are zero (both are proportional to w2). The
            # gradient is exactly zero forever and the experts can never leave.
            #
            # This is what happened to 1p6b-pretrain-v1: 768 routed-expert
            # tensors zero in every checkpoint, 48/48 optimizer moments exactly
            # zero after 15,337 steps, and a routed contribution of norm 0.0.
            # The reference backend was unaffected because its experts are
            # SwiGLU modules made of nn.Linear. deterministic_init would also
            # have covered it -- it walks named_parameters() -- and it was off.
            for w in (module.w1, module.w3, module.w2):
                nn.init.normal_(w, mean=0.0, std=std)

    @torch.no_grad()
    def init_weights_by_name(self, base_seed: int = 0) -> None:
        """Initialise every parameter from a generator seeded by its own *name*.

        Why this exists: `self.apply(init)` walks modules in construction order
        and draws from one global RNG stream, so removing a module (e.g. turning
        Engram off) shifts every subsequent draw. Two ablation arms would then
        differ in the initialisation of *every* layer, not just in the component
        under test -- which would make the comparison meaningless.

        Seeding per parameter name makes `encoder.3.attn.q_proj.weight` identical
        across every configuration that contains it, regardless of what else the
        model does or does not have. That is the property an ablation needs.

        The distributions are unchanged from `_init_weights`; only the source of
        randomness moves.
        """
        import hashlib

        std = self.cfg.initializer_range
        # Which parameters are norms/gates and must start at one, not noise.
        for name, param in self.named_parameters():
            digest = hashlib.sha256(f"{base_seed}:{name}".encode()).digest()
            seed = int.from_bytes(digest[:8], "little") % (2**63 - 1)
            gen = torch.Generator(device="cpu").manual_seed(seed)

            if name.endswith("q_weight") or name.endswith("k_weight"):
                param.fill_(1.0)                      # Engram gate scales
            elif "norm" in name and param.ndim == 1:
                param.fill_(1.0)                      # RMSNorm
            elif name.endswith("attn_sink"):
                param.zero_()
            elif name.endswith(".bias"):
                param.zero_()
            elif ".ffn.router.weight" in name:
                sample = torch.empty(param.shape, dtype=torch.float32)
                sample.normal_(0.0, self.cfg.hidden_size**-0.5, generator=gen)
                param.copy_(sample)
            elif param.ndim >= 2:
                sample = torch.empty(param.shape, dtype=torch.float32)
                sample.normal_(0.0, std, generator=gen)
                param.copy_(sample)
            else:
                param.zero_()

        if self.cfg.tie_word_embeddings:
            # named_parameters() yields a tied weight once, but which name it
            # reports depends on registration order; re-tie explicitly.
            self.lm_head.weight = self.embed.weight

    # ------------------------------------------------------------------ #
    # Encoder
    # ------------------------------------------------------------------ #

    def run_encoder(
        self,
        input_ids: torch.Tensor,
        start_pos: int = 0,
        cache: CEDCache | None = None,
        token_mask: torch.Tensor | None = None,
    ):
        """Returns (H_enc, router_stats, engram_stats)."""
        h = self.embed(input_ids)

        hash_ids = None
        if self.hasher is not None:
            history = cache.engram_history if cache is not None else None
            hash_ids = self.hasher(input_ids, history=history, token_mask=token_mask)
            if cache is not None:
                cache.append_engram_history(self.hasher.compress(input_ids))

        router_stats, engram_stats = [], []
        # `step.layer_index` picks the weights, `step.step_index` picks the cache
        # slot. They are equal under the sequential plan; keeping them separate
        # is what would let a layer run twice without corrupting its own K/V.
        for step in self.encoder_plan:
            i = step.layer_index
            block = self.encoder[i]
            engram = self.engrams[str(i)] if str(i) in self.engrams else None
            if engram is not None:
                h = engram(h, hash_ids[:, :, engram.layer_index, :], token_mask)
                if self.collect_telemetry:
                    engram_stats.append(
                        {"layer": i, **engram.utilization(hash_ids[:, :, engram.layer_index, :])}
                    )
                if cache is not None:
                    cache.cost.engram_calls += 1
            past = cache.enc_kv[step.step_index] if cache is not None else None
            h_in = h
            h, present, route = block(h, self.rope, start_pos=start_pos, past_kv=past)
            if step.residual_scale != 1.0:
                # Scale what this execution ADDED, not its output. Guarded on
                # != 1.0 so the baseline never evaluates `h_in + 1.0 * (h - h_in)`,
                # which is algebraically identity and not bitwise identity.
                h = h_in + step.residual_scale * (h - h_in)
            if cache is not None:
                cache.enc_kv[step.step_index] = present
                cache.cost.encoder_layer_calls += 1
                cache.cost.encoder_tokens += h.shape[0] * h.shape[1]
            if route is not None:
                # Keyed by visit as well as layer: without it the two visits of
                # a looped MoE layer collapse into one set of statistics and
                # a per-visit collapse becomes invisible.
                router_stats.append({
                    "layer": i, "kind": "encoder",
                    "visit": step.visit, "step": step.step_index,
                    **route.stats,
                })
                self._collect_route(route)
        return h, router_stats, engram_stats

    # ------------------------------------------------------------------ #
    # Decoder
    # ------------------------------------------------------------------ #

    def run_decoder(
        self,
        h: torch.Tensor,
        h_enc: torch.Tensor,
        start_pos: int = 0,
        cache: CEDCache | None = None,
        enc_start_pos: int | None = None,
    ):
        """`h` is the decoder input; `h_enc` is the encoder output the global KV
        is projected from.

        These carry separate start positions because truncated prefill runs the
        decoder over a suffix while the global KV must still cover the whole
        prompt -- the decoder's unbounded-range path reads positions 0..t, so
        building it from the suffix only would silently drop the prompt's head.
        """
        cfg = self.cfg
        router_stats = []
        if enc_start_pos is None:
            enc_start_pos = start_pos

        global_kv_all: list[tuple[torch.Tensor, torch.Tensor]] | None = None
        if cfg.ced_enabled:
            if cache is not None:
                cache.cost.global_kv_calls += 1
            enc_normed = self.enc_out_norm(h_enc)
            if cfg.decoder_kv_source == "shared":
                gk, gv = self.global_kv(enc_normed, self.rope, start_pos=enc_start_pos)
                if cache is not None:
                    cache.append_global(gk, gv)
                    gk, gv = cache.global_kv
                global_kv_all = [(gk, gv)] * cfg.decoder_layers
            else:
                global_kv_all = []
                for j, proj in enumerate(self.global_kv):
                    gk, gv = proj(enc_normed, self.rope, start_pos=enc_start_pos)
                    if cache is not None:
                        cache.append_global(gk, gv, layer=j)
                        gk, gv = cache.global_kv_per_layer[j]
                    global_kv_all.append((gk, gv))

        kv_offset = cache.dec_kv_offset if cache is not None else 0
        for step in self.decoder_plan:
            j = step.layer_index
            block = self.decoder[j]
            past = cache.dec_kv[step.step_index] if cache is not None else None
            gkv = global_kv_all[j] if global_kv_all is not None else None
            h, present, route = block(
                h, self.rope, start_pos=start_pos, past_kv=past, global_kv=gkv,
                kv_offset=kv_offset,
            )
            if cache is not None:
                cache.dec_kv[step.step_index] = present
                cache.cost.decoder_layer_calls += 1
                cache.cost.decoder_tokens += h.shape[0] * h.shape[1]
            if route is not None:
                router_stats.append(
                    {"layer": cfg.encoder_layers + j, "kind": "decoder", **route.stats}
                )
                self._collect_route(route)
        if cache is not None:
            cache.cost.decoder_positions += h.shape[1]
        return h, router_stats

    # ------------------------------------------------------------------ #

    def _collect_route(self, route) -> None:
        self._aux_losses.append(route.aux_loss)
        self._z_losses.append(route.z_loss)

    def forward(
        self,
        input_ids: torch.Tensor,
        labels: torch.Tensor | None = None,
        start_pos: int = 0,
        cache: CEDCache | None = None,
        token_mask: torch.Tensor | None = None,
        return_cache: bool = False,
    ) -> ModelOutput:
        self._aux_losses: list[torch.Tensor] = []
        self._z_losses: list[torch.Tensor] = []

        if return_cache and cache is None:
            cache = CEDCache.empty(
                self.cfg, self.encoder_plan.n_steps, self.decoder_plan.n_steps
            )

        h_enc, enc_stats, engram_stats = self.run_encoder(
            input_ids, start_pos=start_pos, cache=cache, token_mask=token_mask
        )
        h, dec_stats = self.run_decoder(h_enc, h_enc, start_pos=start_pos, cache=cache)

        if cache is not None:
            cache.seq_len = start_pos + input_ids.shape[1]

        logits = self.lm_head(self.final_norm(h))

        loss = lm_loss = aux = z = None
        if labels is not None:
            lm_loss = F.cross_entropy(
                logits[:, :-1].reshape(-1, self.cfg.vocab_size).float(),
                labels[:, 1:].reshape(-1),
                ignore_index=-100,
            )
            device = logits.device
            aux = (
                torch.stack(self._aux_losses).mean()
                if self._aux_losses
                else torch.zeros((), device=device)
            )
            z = (
                torch.stack(self._z_losses).mean()
                if self._z_losses
                else torch.zeros((), device=device)
            )
            loss = (
                lm_loss
                + self.cfg.moe.aux_loss_coef * aux
                + self.cfg.moe.z_loss_coef * z
            )

        return ModelOutput(
            logits=logits,
            loss=loss,
            lm_loss=lm_loss,
            aux_loss=aux,
            z_loss=z,
            cache=cache,
            hidden_encoder=h_enc,
            router_stats=enc_stats + dec_stats,
            engram_stats=engram_stats,
        )

    # ------------------------------------------------------------------ #
    # Inference
    # ------------------------------------------------------------------ #

    @property
    def decoder_receptive_field(self) -> int:
        """How many trailing positions the decoder must run over to be exact.

        This is the number that makes CED's prefill/decode asymmetry real, so it
        is worth deriving rather than asserting.

        Decoder layer 0 reads `H_enc`, which is exact at every position. If the
        decoder is run only over the suffix starting at `s`, layer 0 has local
        K/V only for positions >= s, so its attention is correct from
        `s + (W-1)` onward. Each further layer loses another `W-1` positions,
        because layer j's window reaches back `W-1` into layer j-1's outputs.
        Layer j is therefore correct from `s + (j+1)(W-1)`.

        For the final decoder layer to be correct at the last position, and for
        its cached K/V over the trailing window to be correct too (so that later
        decode steps stay exact), we need

            R = W + (D-1)(W-1)

        trailing positions.  Crucially **R does not depend on the prompt length**,
        while the encoder must process all L positions. So as L grows, prefill
        converges to encoder-only work -- which is exactly the "8B active at
        prefill, 16B at decode" asymmetry the model card reports.

        A full-attention decoder (`decoder_window_size == 0`) has unbounded
        receptive field and cannot be truncated at all.
        """
        cfg = self.cfg
        w = cfg.decoder_window_size
        if not cfg.ced_enabled or w <= 0:
            return 2**62  # effectively infinite: no truncation is exact
        return w + (cfg.decoder_layers - 1) * (w - 1)

    def decoder_span_for(self, num_logits: int) -> int:
        """Trailing decoder positions needed to return `num_logits` exact logits.

        The receptive field R buys exactly ONE correct position -- the last.
        Positions earlier in the span are still inside the warm-up region where
        the sliding window is truncated, so their logits are wrong. Each extra
        logit therefore costs one extra position.
        """
        return self.decoder_receptive_field + num_logits - 1

    @torch.no_grad()
    def prefill(
        self,
        input_ids: torch.Tensor,
        num_logits: int = 1,
        full_decoder: bool = False,
    ) -> tuple[torch.Tensor, CEDCache]:
        """Process a prompt and return `(logits, state)`.

        The encoder runs over every position -- it must, because the decoder's
        global KV is projected from `H_enc` at every position. The decoder runs
        only over the trailing positions needed for the requested logits to be
        exact, which is `decoder_span_for(num_logits)` and is **independent of
        prompt length**. That constant-vs-linear split is the CED asymmetry.

        `num_logits`: how many trailing positions to return logits for.
            Generation needs 1; scoring a whole prompt needs `full_decoder`.
        `full_decoder`: run the decoder over the entire prompt and return all
            L logits. Same numbers, more work -- tests use it to prove the
            truncated path agrees with the untruncated one.

        `state.cost` records exactly which modules ran over how many tokens.
        """
        if num_logits < 1:
            raise ValueError("num_logits must be >= 1")
        cache = CEDCache.empty(
            self.cfg, self.encoder_plan.n_steps, self.decoder_plan.n_steps
        )
        cache.cost.reset()
        self._aux_losses, self._z_losses = [], []

        length = input_ids.shape[1]
        cache.cost.positions = length

        h_enc, _, _ = self.run_encoder(input_ids, start_pos=0, cache=cache)

        if full_decoder:
            span, want = length, length
        else:
            want = min(num_logits, length)
            span = min(length, self.decoder_span_for(want))
        start = length - span
        cache.dec_kv_offset = start

        # Decoder input is the suffix; the global KV is still projected from the
        # WHOLE encoder output, because the decoder attends to positions 0..t.
        h, _ = self.run_decoder(
            h_enc[:, start:], h_enc, start_pos=start, cache=cache, enc_start_pos=0
        )
        cache.seq_len = length

        logits = self.lm_head(self.final_norm(h))
        # Only the trailing `want` positions are guaranteed exact; the rest of
        # the span exists solely to warm the sliding windows. Returning them
        # would be handing back silently-wrong numbers.
        return logits[:, -want:], cache

    @torch.no_grad()
    def decode(
        self, tokens: torch.Tensor, state: CEDCache, reset_cost: bool = True
    ) -> tuple[torch.Tensor, CEDCache]:
        """Consume one or more new tokens, returning `(logits, state)`.

        `tokens` is `[B, N]`; N > 1 is accepted (speculative-style block decode).
        The state is mutated in place and also returned, so the call reads as a
        pure state transition at the call site.
        """
        if reset_cost:
            state.cost.reset()
        start_pos = state.seq_len
        state.cost.positions = tokens.shape[1]
        self._aux_losses, self._z_losses = [], []

        h_enc, _, _ = self.run_encoder(tokens, start_pos=start_pos, cache=state)
        h, _ = self.run_decoder(h_enc, h_enc, start_pos=start_pos, cache=state)
        state.seq_len = start_pos + tokens.shape[1]
        return self.lm_head(self.final_norm(h)), state

    @torch.no_grad()
    def decode_step(self, token: torch.Tensor, cache: CEDCache) -> torch.Tensor:
        """Backwards-compatible single-step alias for `decode`."""
        return self.decode(token, cache)[0]

    @torch.no_grad()
    def generate(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 32,
        temperature: float = 1.0,
        top_k: int | None = None,
        eos_token_id: int | None = None,
    ) -> torch.Tensor:
        self.eval()
        if input_ids.shape[1] > self.cfg.max_sequence_length:
            raise ValueError(
                f"prompt of {input_ids.shape[1]} tokens exceeds "
                f"max_sequence_length={self.cfg.max_sequence_length}"
            )
        logits, cache = self.prefill(input_ids)
        out = input_ids
        for _ in range(max_new_tokens):
            # Check BEFORE emitting: appending first would return a sequence
            # longer than the model can process, which then cannot be fed back
            # in or resumed from.
            if cache.seq_len >= self.cfg.max_sequence_length:
                break
            next_logits = logits[:, -1].float()
            if temperature <= 0:
                nxt = next_logits.argmax(-1, keepdim=True)
            else:
                next_logits = next_logits / temperature
                if top_k:
                    kth = next_logits.topk(min(top_k, next_logits.shape[-1]), dim=-1)[0][:, -1:]
                    next_logits = next_logits.masked_fill(next_logits < kth, float("-inf"))
                nxt = torch.multinomial(next_logits.softmax(-1), 1)
            out = torch.cat([out, nxt], dim=1)
            if eos_token_id is not None and bool((nxt == eos_token_id).all()):
                break
            logits, cache = self.decode(nxt, cache)
        return out

    # ------------------------------------------------------------------ #
    # Introspection
    # ------------------------------------------------------------------ #

    def parameter_report(self) -> dict:
        """Real parameter counts by role, plus an active-per-token estimate."""
        cfg = self.cfg
        buckets = {
            "embeddings": 0,
            "attention": 0,
            "global_kv": 0,
            "routed_experts": 0,
            "shared_experts": 0,
            "router": 0,
            "engram_tables": 0,
            "engram_other": 0,
            "dense_ffn": 0,
            "norms": 0,
            "lm_head": 0,
            "other": 0,
        }
        seen: set[int] = set()
        for name, p in self.named_parameters():
            if id(p) in seen:  # tied weights counted once
                continue
            seen.add(id(p))
            n = p.numel()
            if name.startswith("embed."):
                key = "embeddings"
            elif name.startswith("lm_head"):
                key = "lm_head"
            elif ".attn." in name:
                key = "attention"
            elif name.startswith("global_kv"):
                key = "global_kv"
            elif ".ffn.experts." in name or ".ffn.grouped." in name:
                # Both backends hold the same weights; only the layout differs.
                # `.ffn.experts.<i>.*` is the reference bank, `.ffn.grouped.w*`
                # the stacked one. Missing the second counted 1.21B of routed
                # parameters as dense, which tripled active_per_token and so
                # inflated MFU 3.6x and wrote the wrong active_parameter_count
                # into every checkpoint the capability lab reads.
                key = "routed_experts"
            elif ".ffn.shared_expert." in name:
                key = "shared_experts"
            elif ".ffn.router." in name:
                key = "router"
            elif ".ffn.ffn." in name:
                key = "dense_ffn"
            elif name.startswith("engrams.") and ".table." in name:
                key = "engram_tables"
            elif name.startswith("engrams."):
                key = "engram_other"
            elif "norm" in name:
                key = "norms"
            else:
                key = "other"
            buckets[key] += n

        total = sum(buckets.values())
        dense = total - buckets["routed_experts"] - buckets["engram_tables"]

        # Active per token: every routed expert param is replaced by top_k/N of it,
        # and the Engram table by the handful of rows actually read.
        if cfg.moe.enabled and cfg.moe.num_routed_experts:
            active_routed = buckets["routed_experts"] * cfg.moe.top_k / cfg.moe.num_routed_experts
        else:
            active_routed = buckets["routed_experts"]
        active_engram = 0
        if self.engram_layout is not None:
            active_engram = (
                self.engram_layout.n_hash_cols
                * self.engram_layout.embedding_dim
                * len(self.engram_layout.layer_ids)
            )
        active = (
            buckets["embeddings"]
            + buckets["attention"]
            + buckets["global_kv"]
            + buckets["shared_experts"]
            + buckets["router"]
            + buckets["dense_ffn"]
            + buckets["norms"]
            + buckets["lm_head"]
            + buckets["other"]
            + buckets["engram_other"]
            + active_routed
            + active_engram
        )

        report = dict(buckets)
        report.update(
            {
                "total": total,
                "dense_total": dense,
                "active_per_token": int(active),
                "engram_rows": (
                    list(self.engram_layout.num_embeddings) if self.engram_layout else []
                ),
                "engram_hash_cols": (
                    self.engram_layout.n_hash_cols if self.engram_layout else 0
                ),
                "engram_compressed_vocab": (
                    self.engram_layout.compressed_vocab_size if self.engram_layout else 0
                ),
            }
        )
        return report

    def num_parameters(self) -> int:
        seen, total = set(), 0
        for p in self.parameters():
            if id(p) not in seen:
                seen.add(id(p))
                total += p.numel()
        return total