model.py

"""
Quazimoto-LM — a language model whose channel-mixing block is a bank of
coupled phase oscillators (Kuramoto dynamics) arranged in concentric rings.

Backbone: standard causal transformer (token mixing = attention, tied LM head).
Novelty:  the per-layer MLP is replaced by a QuazimotoBlock that maps the hidden
          state to N oscillator phases on rings [4, 8, 16, 32, ...], runs a few
          differentiable Euler steps of structured Kuramoto coupling with a
          learnable frustration alpha, then reads out [cos, sin] of the phases.

Set ring coupling to dense + alpha=0 and it degenerates toward AKOrN-style
oscillatory neurons; the ring structure + hierarchy-to-center is the Quazimoto bit.
"""

from dataclasses import dataclass, field
import math
import os
import torch
import torch.nn as nn
import torch.nn.functional as F

_PKG_DIR = os.path.dirname(os.path.abspath(__file__))

# family traits (DERF soft_clamp, RMSNorm, and the opt-in refinement/aux blocks)
from family import (soft_clamp, RMSNorm, HRMRefinementBlock, MoESwiGLU,
                    MTPHead, JEPAPredictorBlock, PhaseAttentionRing, EngramRing,
                    RingControllerBank, RingSpecialists)
import instrument as _viz   # live-visualizer capture hooks (no-op unless a recorder is active)


@dataclass
class QuazimotoConfig:
    vocab_size: int = 32000
    n_layer: int = 10
    n_head: int = 12
    d_model: int = 768
    block_size: int = 1024
    dropout: float = 0.0
    # Quazimoto oscillator bank
    ring_sizes: tuple = (4, 8, 16, 32, 64, 128, 256)  # concentric rings
    osc_steps: int = 4          # differentiable Kuramoto Euler steps
    osc_dt: float = 0.5
    readout_mult: int = 3       # MLP expansion on the [cos,sin] readout  (~98M params)
    # Mandelbrot phase seeding: seed the N oscillator phases from each token's
    # fractal orbit angles (frozen, parameter-free), added through a zero-init gate.
    use_fractal_phase_seed: bool = False
    # family-trait discipline
    osc_bound: float = 10.0     # DERF soft_clamp bound on block output (BPTT-stability)
    gate_init_open: float = 0.9 # tanh-gate opening at init (Quazimoto is the main mixer,
                                #   so it starts OPEN; set 0.0 for a pure no-op refinement)
    # interstitial rings: between each pair of oscillator rings, insert a phase-attention
    # ring + an engram ring that ABSORB info and INJECT it into the two neighbor rings.
    use_rings: bool = False
    ring_attn_heads: int = 4
    ring_attn_head_dim: int = 16
    ring_engram_compress: int = 32
    ring_engram_heads: int = 2
    ring_engram_table: int = 2048
    ring_engram_ngram: int = 3
    # per-ring memory specialists: a small MoE of `ring_n_specialists` mini experts
    # PER oscillator ring, each holding test-time-writable input/output stores that
    # act as an addressable context memory at inference (see family.RingSpecialists).
    use_ring_specialists: bool = False
    ring_n_specialists: int = 7
    ring_spec_key_dim: int = 32
    ring_spec_slot_dim: int = 64
    ring_spec_top_k: int = 2
    ring_spec_write_lr: float = 0.1
    # per-ring self-organizing controllers (~1M total: one tiny net per oscillator
    # ring, shared across layers, weights self-optimize by a predictive/free-energy rule)
    use_ring_controllers: bool = False
    ring_ctrl_feat: int = 384        # fast-weight predictor width (dominates the ~150k/ring)
    ring_ctrl_local_lr: float = 0.01 # delta-rule (surprise-minimization) step size
    # opt-in family trait modules (all safe no-ops at init)
    use_hrm: bool = False
    hrm_steps: int = 3
    hrm_dim: int = 256
    hrm_gate_init: float = 0.1   # HRM gates start OPEN (random initial state needs a path out)
    use_moe: bool = False
    moe_intermediate: int = 768
    moe_n_routed: int = 4
    moe_n_shared: int = 1
    moe_top_k: int = 2
    use_mtp: bool = False
    mtp_layers: int = 4   # depth of MTP draft heads (predict +1..+mtp_layers); used for spec-decode
    mtp_loss_weight: float = 0.3
    use_jepa: bool = False
    jepa_pred_dim: int = 256
    jepa_horizon: int = 1
    jepa_loss_weight: float = 0.1
    # ---- attention: ported from model_v2.MLADerfXSAAttention ----
    head_dim: int = 64
    qk_rope_head_dim: int = 32       # partial RoPE: rope part of each head
    nope_head_dim: int = 32          #               no-pos part (sum == head_dim)
    num_key_value_heads: int = 4     # GQA (== n_head for full MHA)
    q_lora_rank: int = 384           # MLA low-rank q projection
    o_lora_rank: int = 384           # MLA low-rank output projection
    use_qk_norm: bool = True         # per-head RMSNorm on Q/K before RoPE
    use_derf: bool = False           # erf attention instead of softmax (ablate)
    use_xsa: bool = False            # value-subspace removal (ablate)
    rope_theta: float = 10000.0
    max_position_embeddings: int = 4096
    rms_norm_eps: float = 1e-6
    initializer_range: float = 0.02
    # ---- model-level v2 features ----
    zloss_coef: float = 1e-4         # z-loss on lm logits
    use_value_embed: bool = False    # per-block value-embedding residual (zero-init gate)
    use_hyper_connections: bool = False
    hc_mult: int = 2

    @property
    def n_osc(self):
        return sum(self.ring_sizes)


def _ring_index(ring_sizes):
    """Return a LongTensor of length N mapping each oscillator -> its ring id."""
    idx = []
    for r, n in enumerate(ring_sizes):
        idx += [r] * n
    return torch.tensor(idx, dtype=torch.long)


class QuazimotoBlock(nn.Module):
    """Oscillatory channel-mixing block: hidden state -> ring phases -> Kuramoto -> readout."""

    def __init__(self, cfg: QuazimotoConfig):
        super().__init__()
        self.cfg = cfg
        N = cfg.n_osc
        R = len(cfg.ring_sizes)
        self.N, self.R = N, R

        self.norm = RMSNorm(cfg.d_model)
        # project hidden -> initial phases and natural frequencies
        self.to_theta = nn.Linear(cfg.d_model, N)
        self.to_omega = nn.Linear(cfg.d_model, N)

        # structured, LEARNABLE coupling: an RxR gain between ring blocks.
        # Block-constant coupling lets us use per-ring mean fields (order
        # parameters) -> O(N*R) instead of the O(N^2) pairwise sum.
        rid = _ring_index(cfg.ring_sizes)
        self.register_buffer("ring_id", rid)
        onehot = F.one_hot(rid, R).float()       # [N,R] ring membership
        self.register_buffer("onehot", onehot)
        self.block_gain = nn.Parameter(self._init_gain(R))
        self.center_phase = nn.Parameter(torch.zeros(()))  # the rotating center ball
        self.center_omega = nn.Parameter(torch.tensor(0.3))
        self.center_gain = nn.Parameter(torch.full((R,), 0.5))  # ring -> center pull
        self.alpha = nn.Parameter(torch.zeros(R))              # per-ring frustration

        # readout MLP on [cos theta, sin theta]
        hidden = cfg.readout_mult * cfg.d_model
        self.readout = nn.Sequential(
            nn.Linear(2 * N, hidden),
            nn.GELU(),
            nn.Linear(hidden, cfg.d_model),
        )
        self.drop = nn.Dropout(cfg.dropout)

        # family-trait gate: out = readout * tanh(gate). gate_init_open seeds it so the
        # block is a warm-startable signal whose magnitude the optimizer controls; the
        # content weights are nonzero so the gate always receives gradient.
        go = math.atanh(min(cfg.gate_init_open, 0.9)) if cfg.gate_init_open > 0 else 0.0
        self.gate = nn.Parameter(torch.full((1,), go))

        # Mandelbrot phase-seed gate: theta <- to_theta(h) + tanh(seed_gate) * orbit.
        # Zero-init -> no-op at start; opens only if the fractal prior helps.
        if cfg.use_fractal_phase_seed:
            self.phase_seed_gate = nn.Parameter(torch.zeros(1))

        # std=0.02 init discipline (matches the family's linear-layer init)
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, std=0.02)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

        # interstitial rings (instantiated AFTER the std-init loop). Oscillators are
        # ordered ring-by-ring, so rings r and r+1 are the contiguous slice
        # [off_r, off_{r+2}); the slot between them injects into exactly that block.
        self.use_rings = cfg.use_rings
        self.slot_bounds = []
        if self.use_rings:
            off = [0]
            for n in cfg.ring_sizes:
                off.append(off[-1] + n)
            self.slot_bounds = [(off[r], off[r + 2]) for r in range(R - 1)]
            self.attn_rings = nn.ModuleList([
                PhaseAttentionRing(e - s, cfg.ring_attn_heads, cfg.ring_attn_head_dim, cfg.osc_bound)
                for (s, e) in self.slot_bounds])
            self.engram_rings = nn.ModuleList([
                EngramRing(cfg.d_model, e - s, cfg.ring_engram_compress, cfg.ring_engram_heads,
                           cfg.ring_engram_table, cfg.ring_engram_ngram, cfg.osc_bound)
                for (s, e) in self.slot_bounds])

        # one memory-specialist MoE bank PER oscillator ring (injects into that ring)
        self.use_ring_specialists = cfg.use_ring_specialists
        if self.use_ring_specialists:
            self.specialists = nn.ModuleList([
                RingSpecialists(n, cfg.d_model, cfg.ring_n_specialists, cfg.ring_spec_key_dim,
                                cfg.ring_spec_slot_dim, cfg.ring_spec_top_k,
                                cfg.ring_spec_write_lr, cfg.osc_bound)
                for n in cfg.ring_sizes])

    @staticmethod
    def _init_gain(R):
        # within-ring coupling strong on the diagonal, weak neighbor coupling
        g = torch.zeros(R, R)
        for i in range(R):
            g[i, i] = 1.0
            if i + 1 < R:
                g[i, i + 1] = g[i + 1, i] = 0.3
        return g

    def forward(self, x, ring_ctl=None, phase_seed=None):
        cfg = self.cfg
        B, T, _ = x.shape
        h = self.norm(x)

        theta = self.to_theta(h)                 # [B,T,N]  initial phases
        if phase_seed is not None:               # Mandelbrot orbit phase prior (gated, no-op at init)
            theta = theta + torch.tanh(self.phase_seed_gate) * phase_seed
        omega = torch.tanh(self.to_omega(h))     # [B,T,N]  natural frequencies (bounded)

        # per-ring self-organizing controller: observe each ring (order-parameter
        # magnitude/phase + mean freq), get no-op-at-init modulations of that ring's
        # coupling / frustration / center-pull / injection.
        d_coup = d_alpha = d_center = d_inj = None
        if ring_ctl is not None:
            with torch.no_grad():
                c0, s0 = torch.cos(theta), torch.sin(theta)
                cnt = self.onehot.sum(0).clamp(min=1.0)       # [R] ring sizes
                zc = (c0 @ self.onehot) / cnt                 # [B,T,R]
                zs = (s0 @ self.onehot) / cnt
                mag = torch.sqrt(zc ** 2 + zs ** 2 + 1e-8)
                om = (omega @ self.onehot) / cnt
                obs = torch.stack([mag.mean((0, 1)), (zc / mag).mean((0, 1)),
                                   (zs / mag).mean((0, 1)), om.mean((0, 1))], dim=-1)  # [R,4]
            ctrl = ring_ctl(obs)                              # [R,4] (grad -> enc/dec only)
            d_coup = torch.tanh(ctrl[:, 0])                   # ring coupling scale (1 + .)
            d_alpha = 0.5 * torch.tanh(ctrl[:, 1])            # frustration shift
            d_center = torch.tanh(ctrl[:, 2])                 # center-pull scale (1 + .)
            d_inj = torch.tanh(ctrl[:, 3])                    # injection scale (1 + .)

        # interstitial rings absorb info ONCE (from the initial phases / hidden) and
        # inject a constant drive into their two neighbor oscillator rings' phase update.
        inject = 0.0
        if self.use_rings:
            inject = torch.zeros_like(theta)
            for (s, e), ar, er in zip(self.slot_bounds, self.attn_rings, self.engram_rings):
                contrib = ar(theta[..., s:e]) + er(h)         # [B,T,e-s]
                inject = inject + F.pad(contrib, (s, self.N - e))  # place at [s:e], sum overlaps
            if d_inj is not None:
                inject = inject * (1.0 + d_inj)[self.ring_id]      # controller scales the drive

        # memory specialists: each per-ring bank injects retrieved store_out into its
        # ring's oscillators. Rings are contiguous and tile [0,N), so cat == full drive.
        if self.use_ring_specialists:
            spec_inj = torch.cat([bank(h) for bank in self.specialists], dim=-1)  # [B,T,N]
            inject = inject + spec_inj

        rid = self.ring_id                        # [N]
        gain = self.block_gain if d_coup is None else self.block_gain * (1.0 + d_coup).unsqueeze(1)
        G = gain[rid]                             # [N,R]  per-oscillator ring gains
        alpha_vec = self.alpha if d_alpha is None else self.alpha + d_alpha
        alpha_i = alpha_vec[rid]                  # [N]
        center = self.center_phase                # scalar phase, advances in time
        cpull = self.center_gain if d_center is None else self.center_gain * (1.0 + d_center)
        center_pull = cpull[rid]                  # [N]
        invN = 1.0 / self.N

        for s in range(cfg.osc_steps):
            ctr = center + self.center_omega * (s * cfg.osc_dt)
            c, sn = torch.cos(theta), torch.sin(theta)            # [B,T,N]
            # per-ring summed mean fields, then routed back through gains G
            zc = (c @ self.onehot) @ G.t()                        # [B,T,N]
            zs = (sn @ self.onehot) @ G.t()                       # [B,T,N]
            A = alpha_i - theta                                   # [B,T,N]
            # (1/N) * sum_r G_ir * Im( e^{i(alpha_i - theta_i)} * Zsum_r )
            coupling = invN * (torch.sin(A) * zc + torch.cos(A) * zs)
            to_center = center_pull * torch.sin(ctr - theta + alpha_i)            # [B,T,N]
            theta = theta + cfg.osc_dt * (omega + coupling + to_center + inject)  # ring drive

        feat = torch.cat([torch.cos(theta), torch.sin(theta)], dim=-1)  # [B,T,2N]
        out = self.readout(feat) * torch.tanh(self.gate)
        out = self.drop(soft_clamp(out, self.cfg.osc_bound))           # bounded, live grad

        rec = _viz.get_rec()
        if rec is not None and rec.enabled:                            # live-viz capture
            th = theta[0, -1]                                          # last-token phases [N]
            cnt = self.onehot.sum(0).clamp(min=1.0)
            zc = (torch.cos(th) @ self.onehot) / cnt                   # per-ring order param
            zs = (torch.sin(th) @ self.onehot) / cnt
            R = torch.sqrt(zc ** 2 + zs ** 2 + 1e-9)
            psi = torch.atan2(zs, zc)
            rec.log_ring(R.tolist(), psi.tolist(), th.tolist())
            rec.log_quaz_norm(out[0, -1].norm().item())
            # Mandelbrot phase seed: this layer's injection gate + (for the chosen
            # layer) the token's orbit-angle seed -- the initial phases the fractal
            # prior hands the oscillator bank, before any Kuramoto step.
            if cfg.use_fractal_phase_seed and phase_seed is not None:
                rec.log_fractal(phase_seed[0, -1].tolist(),
                                float(torch.tanh(self.phase_seed_gate)))
            rec.flush_spec()                                           # group this layer's routes
        return out


class RotaryEmbedding(nn.Module):
    """RoPE for the rope partition of Q/K (qk_rope_head_dim dims). Ported from v2."""

    def __init__(self, dim, max_positions=4096, theta=10000.0):
        super().__init__()
        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
        t = torch.arange(max_positions).float()
        freqs = torch.outer(t, inv_freq)
        self.register_buffer("cos_cache", freqs.cos(), persistent=False)
        self.register_buffer("sin_cache", freqs.sin(), persistent=False)

    def forward(self, x, position_ids):
        cos = self.cos_cache[position_ids].unsqueeze(1)
        sin = self.sin_cache[position_ids].unsqueeze(1)
        d = cos.shape[-1]
        x1, x2 = x[..., :d], x[..., d:]
        return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)


class MLADerfXSAAttention(nn.Module):
    """Family attention, ported from model_v2.MLADerfXSAAttention: MLA low-rank
    q/o projections, partial RoPE (nope+rope split), QK-Norm, optional DERF erf
    attention, optional XSA value-subspace removal, GQA. Supports an optional KV
    cache for incremental decoding (cache holds post-RoPE, pre-GQA-expansion k/v
    of shape [B, num_kv_heads, S, head_dim])."""

    def __init__(self, cfg: QuazimotoConfig):
        super().__init__()
        self.num_heads = cfg.n_head
        self.num_kv_heads = cfg.num_key_value_heads
        self.head_dim = cfg.head_dim
        self.nope_head_dim = cfg.nope_head_dim
        self.use_qk_norm = cfg.use_qk_norm
        self.use_derf = cfg.use_derf
        self.use_xsa = cfg.use_xsa
        self.dropout_p = cfg.dropout
        self.kv_groups = self.num_heads // self.num_kv_heads
        assert self.nope_head_dim + cfg.qk_rope_head_dim == self.head_dim, \
            "nope_head_dim + qk_rope_head_dim must equal head_dim"

        self.q_a_proj = nn.Linear(cfg.d_model, cfg.q_lora_rank, bias=False)
        self.q_a_norm = RMSNorm(cfg.q_lora_rank, cfg.rms_norm_eps)
        self.q_b_proj = nn.Linear(cfg.q_lora_rank, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(cfg.d_model, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(cfg.d_model, self.num_kv_heads * self.head_dim, bias=False)
        self.o_a_proj = nn.Linear(self.num_heads * self.head_dim, cfg.o_lora_rank, bias=False)
        self.o_b_proj = nn.Linear(cfg.o_lora_rank, cfg.d_model, bias=False)

        if self.use_qk_norm:
            self.q_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps)
            self.k_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps)
        self.rope = RotaryEmbedding(cfg.qk_rope_head_dim, cfg.max_position_embeddings,
                                    cfg.rope_theta)
        if self.use_derf:
            self.derf_alpha = nn.Parameter(torch.ones(self.num_heads))
            self.derf_bias = nn.Parameter(torch.zeros(self.num_heads))
            self.derf_gamma = nn.Parameter(torch.ones(self.num_heads))
        for m in (self.q_a_proj, self.q_b_proj, self.k_proj, self.v_proj,
                  self.o_a_proj, self.o_b_proj):
            nn.init.normal_(m.weight, std=cfg.initializer_range)

    def forward(self, x, position_ids, past_kv=None, use_cache=False):
        B, S, _ = x.shape
        q = self.q_b_proj(self.q_a_norm(self.q_a_proj(x)))
        q = q.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)
        if self.use_qk_norm:
            q, k = self.q_norm(q), self.k_norm(k)

        d = self.nope_head_dim
        q = torch.cat([q[..., :d], self.rope(q[..., d:], position_ids)], dim=-1)
        k = torch.cat([k[..., :d], self.rope(k[..., d:], position_ids)], dim=-1)

        # KV cache: prepend previously-seen post-RoPE k/v (pre-GQA-expansion so
        # the cache stays GQA-compact), then this step's tokens become the tail.
        if past_kv is not None:
            past_k, past_v = past_kv
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)
        present = (k, v) if use_cache else None
        L = k.size(2)                                  # total keys (cache + current)

        if self.kv_groups > 1:
            k = k.unsqueeze(2).expand(-1, -1, self.kv_groups, -1, -1).reshape(
                B, self.num_heads, L, self.head_dim)
            v = v.unsqueeze(2).expand(-1, -1, self.kv_groups, -1, -1).reshape(
                B, self.num_heads, L, self.head_dim)

        # Causal mask over the [S queries x L keys] block. When there is no cache
        # and S == L this is the plain lower triangle; with a cache the S new
        # queries sit at absolute positions [L-S, L) and may attend all keys <=
        # their own position. allowed[i,j] = j <= (L - S + i).
        offset = L - S
        if self.use_derf:
            qpos = torch.arange(S, device=x.device).view(S, 1) + offset
            kpos = torch.arange(L, device=x.device).view(1, L)
            is_masked = (kpos > qpos).unsqueeze(0).unsqueeze(0)   # [1,1,S,L]
            scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
            safe = scores.masked_fill(is_masked, -10000.0)
            a, b, g = (self.derf_alpha.view(1, -1, 1, 1), self.derf_bias.view(1, -1, 1, 1),
                       self.derf_gamma.view(1, -1, 1, 1))
            w = g * torch.erf(a * safe + b)
            w = (w + g) / 2.0
            w = w.masked_fill(is_masked, 0.0)
            w = w / (w.sum(-1, keepdim=True) + 1e-8)
            if self.dropout_p > 0 and self.training:
                w = F.dropout(w, p=self.dropout_p)
            y = torch.matmul(w, v)
        else:
            if offset == 0:
                attn_mask, causal = None, True
            else:
                qpos = torch.arange(S, device=x.device).view(S, 1) + offset
                kpos = torch.arange(L, device=x.device).view(1, L)
                attn_mask = (kpos <= qpos).unsqueeze(0).unsqueeze(0)  # bool: True=keep
                causal = False
            y = F.scaled_dot_product_attention(
                q.contiguous(), k.contiguous(), v.contiguous(), attn_mask=attn_mask,
                is_causal=causal, dropout_p=self.dropout_p if self.training else 0.0)

        rec = _viz.get_rec()
        if rec is not None and rec.enabled:                  # last-query attention, mean over heads
            sc = (q[:, :, -1:] @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)  # [B,H,1,L]
            w = torch.softmax(sc, dim=-1).mean(1)[0, 0]      # [L]
            rec.log_attn(getattr(self, "layer_idx", 0), w.tolist())

        if self.use_xsa:
            # remove each query's component along its OWN value direction; with a
            # cache the queries are the last S of the L cached value positions.
            vq = v[:, :, -S:, :]
            vn = vq / (vq.norm(dim=-1, keepdim=True) + 1e-8)
            y = y - (y * vn).sum(-1, keepdim=True) * vn

        y = y.transpose(1, 2).contiguous().view(B, S, self.num_heads * self.head_dim)
        out = self.o_b_proj(self.o_a_proj(y))
        return (out, present) if use_cache else out


class HyperConnectionLayer(nn.Module):
    """Softmax pre-mix / post-distribute over hc_mult residual streams (v2)."""

    def __init__(self, hc_mult):
        super().__init__()
        self.pre_weight = nn.Parameter(torch.linspace(0.5, -0.5, hc_mult) / max(hc_mult, 1))
        self.post_weight = nn.Parameter(torch.linspace(-0.5, 0.5, hc_mult) / max(hc_mult, 1))

    def pre_op(self, copies):
        w = F.softmax(self.pre_weight, dim=0)
        return (copies * w.view(1, -1, 1, 1)).sum(dim=1)

    def post_op(self, copies, delta):
        w = F.softmax(self.post_weight, dim=0)
        return copies + delta.unsqueeze(1) * w.view(1, -1, 1, 1)


class HCOutputMix(nn.Module):
    """Learned softmax mix over hc_mult streams at the output (init == mean)."""

    def __init__(self, hc_mult):
        super().__init__()
        self.weight = nn.Parameter(torch.zeros(hc_mult))

    def forward(self, copies):
        w = F.softmax(self.weight, dim=0)
        return (copies * w.view(1, -1, 1, 1)).sum(dim=1)


class Layer(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.use_hc = cfg.use_hyper_connections
        self.use_value_embed = cfg.use_value_embed
        self.attn_norm = RMSNorm(cfg.d_model, cfg.rms_norm_eps)
        self.attn = MLADerfXSAAttention(cfg)
        self.quaz = QuazimotoBlock(cfg)              # the oscillator FFN sub-layer
        if self.use_hc:
            self.hc_attn = HyperConnectionLayer(cfg.hc_mult)
            self.hc_ffn = HyperConnectionLayer(cfg.hc_mult)
        if self.use_value_embed:
            self.ve_gate = nn.Parameter(torch.zeros(1))   # zero-init -> no-op

    def forward(self, x, position_ids, token_embed=None, ring_ctl=None,
                past_kv=None, use_cache=False, phase_seed=None):
        h = self.hc_attn.pre_op(x) if self.use_hc else x
        if self.use_value_embed and token_embed is not None:
            h = h + torch.tanh(self.ve_gate) * token_embed
        attn_out = self.attn(self.attn_norm(h), position_ids,
                             past_kv=past_kv, use_cache=use_cache)
        present = None
        if use_cache:
            attn_out, present = attn_out
        if self.use_hc:
            x = self.hc_attn.post_op(x, attn_out)
            h = self.hc_ffn.pre_op(x)
        else:
            h = h + attn_out
        ffn_out = self.quaz(h, ring_ctl, phase_seed)  # QuazimotoBlock norms internally
        if self.use_hc:
            x = self.hc_ffn.post_op(x, ffn_out)
        else:
            x = h + ffn_out
        return (x, present) if use_cache else x


class QuazimotoLM(nn.Module):
    def __init__(self, cfg: QuazimotoConfig):
        super().__init__()
        self.cfg = cfg
        self.tok = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.drop = nn.Dropout(cfg.dropout)        # positions come from RoPE (no abs embed)
        # frozen Mandelbrot orbit-angle phase table [vocab, n_osc]; non-persistent
        # (deterministic, regenerated on load -> not stored in checkpoints).
        if cfg.use_fractal_phase_seed:
            from fractal import mandelbrot_phase_table, load_phase_table
            pt, mode = load_phase_table(cfg.vocab_size, cfg.n_osc,
                                        os.path.join(_PKG_DIR, "fractal_phase.pt"))
            if pt is None:
                pt = mandelbrot_phase_table(cfg.vocab_size, cfg.n_osc)
                print("  [fractal] flat Halton seed (run build_fractal_table.py for hierarchical)")
            else:
                print(f"  [fractal] loaded {mode} phase table")
            self.register_buffer("fractal_phase", pt, persistent=False)
        self.layers = nn.ModuleList([Layer(cfg) for _ in range(cfg.n_layer)])
        for i, layer in enumerate(self.layers):
            layer.attn.layer_idx = i                 # for live-viz attention capture
        self.norm_f = RMSNorm(cfg.d_model, cfg.rms_norm_eps)
        self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        self.head.weight = self.tok.weight  # weight tying
        self.logit_scale = nn.Parameter(torch.tensor(1.0))  # family trait: learned logit temp
        self.use_value_embed = cfg.use_value_embed
        self.hc_out_mix = HCOutputMix(cfg.hc_mult) if cfg.use_hyper_connections else None
        self.apply(self._init)
        # engram rings manage their own inits (frozen compressor, std-0.01 tables,
        # DERF bias -4); re-apply them after the global init clobbers those.
        for m in self.modules():
            if isinstance(m, (EngramRing, RingSpecialists)):
                m.family_reinit()

        # opt-in trait modules -- instantiated AFTER self.apply so their internal
        # zero-inits (MoE down-proj, MTP proj) and gate inits survive the global init.
        # HRM/MoE refine the trunk after the stack (the family's "top membrane
        # refinement" point); MTP/JEPA are train-time aux heads on the trunk rep.
        self.hrm = (HRMRefinementBlock(cfg.d_model, cfg.hrm_dim, cfg.hrm_steps,
                                       gate_init_open=cfg.hrm_gate_init)
                    if cfg.use_hrm else None)
        self.moe = MoESwiGLU(cfg.d_model, cfg.moe_intermediate, cfg.moe_n_routed,
                             cfg.moe_n_shared, cfg.moe_top_k) if cfg.use_moe else None
        self.mtp_heads = (nn.ModuleList([MTPHead(cfg.d_model) for _ in range(cfg.mtp_layers)])
                          if cfg.use_mtp else None)
        self.jepa = (JEPAPredictorBlock(cfg.d_model, cfg.jepa_pred_dim, cfg.jepa_horizon)
                     if cfg.use_jepa else None)
        self.ring_bank = (RingControllerBank(len(cfg.ring_sizes), 4, cfg.ring_ctrl_feat,
                                             cfg.ring_ctrl_local_lr)
                          if cfg.use_ring_controllers else None)

        # exact kwargs to rebuild on resume (family `family_config` convention)
        self.family_config = {k: getattr(cfg, k) for k in (
            "vocab_size", "n_layer", "n_head", "d_model", "block_size",
            "ring_sizes", "osc_steps", "osc_dt", "readout_mult", "osc_bound", "gate_init_open",
            "use_hrm", "hrm_steps", "hrm_dim", "hrm_gate_init",
            "use_moe", "moe_intermediate", "moe_n_routed",
            "moe_n_shared", "moe_top_k", "use_mtp", "mtp_layers", "mtp_loss_weight",
            "use_jepa", "jepa_pred_dim", "jepa_horizon", "jepa_loss_weight",
            "head_dim", "qk_rope_head_dim", "nope_head_dim", "num_key_value_heads",
            "q_lora_rank", "o_lora_rank", "use_qk_norm", "use_derf", "use_xsa",
            "rope_theta", "max_position_embeddings", "rms_norm_eps", "initializer_range",
            "zloss_coef", "use_value_embed", "use_hyper_connections", "hc_mult",
            "use_rings", "ring_attn_heads", "ring_attn_head_dim", "ring_engram_compress",
            "ring_engram_heads", "ring_engram_table", "ring_engram_ngram",
            "use_ring_controllers", "ring_ctrl_feat", "ring_ctrl_local_lr",
            "use_ring_specialists", "ring_n_specialists", "ring_spec_key_dim",
            "ring_spec_slot_dim", "ring_spec_top_k", "ring_spec_write_lr",
            "use_fractal_phase_seed")}
        n = sum(p.numel() for p in self.parameters())
        print(f"QuazimotoLM: {n/1e6:.1f}M params | {cfg.n_osc} oscillators/token "
              f"over {len(cfg.ring_sizes)} rings {cfg.ring_sizes} | traits: "
              f"{self.active_traits()}")

    def active_traits(self):
        t = []
        if self.cfg.use_hrm: t.append("HRM")
        if self.cfg.use_moe: t.append("MoE")
        if self.cfg.use_mtp: t.append("MTP")
        if self.cfg.use_jepa: t.append("JEPA")
        if self.cfg.use_rings: t.append("Rings")
        if self.cfg.use_ring_controllers: t.append("Controllers")
        if self.cfg.use_ring_specialists:
            t.append(f"Specialists({self.cfg.ring_n_specialists}/ring)")
        if self.cfg.use_fractal_phase_seed: t.append("FractalSeed")
        return ",".join(t) or "none"

    def reset_ring_memory(self):
        """Zero every RingSpecialists store. Call between independent prompts so
        the test-time context memory does not carry over from a previous sequence."""
        for m in self.modules():
            if isinstance(m, RingSpecialists):
                m.reset_memory()

    def set_ring_memory_writing(self, enabled: bool):
        """Toggle online writes to the specialist stores (e.g. freeze the memory
        during evaluation, or disable it to ablate the test-time-write behavior)."""
        for m in self.modules():
            if isinstance(m, RingSpecialists):
                m.write_enabled = enabled

    def _init(self, m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, std=0.02)

    def _compute_trunk(self, idx, past_key_values=None, use_cache=False):
        """Shared backbone: tokens -> layers -> trunk refinements -> RMSNorm trunk.
        Returns (trunk, presents). Used by both forward() and forward_drafts()."""
        B, T = idx.shape
        # with a cache, the new tokens sit at absolute positions [past_len, past_len+T)
        past_len = past_key_values[0][0].size(2) if past_key_values else 0
        position_ids = (torch.arange(past_len, past_len + T, device=idx.device)
                        .unsqueeze(0).expand(B, -1))
        x = self.drop(self.tok(idx))
        token_embed = x if self.use_value_embed else None
        phase_seed = self.fractal_phase[idx] if self.cfg.use_fractal_phase_seed else None
        if self.cfg.use_hyper_connections:
            x = x.unsqueeze(1).expand(-1, self.cfg.hc_mult, -1, -1).clone()  # [B,M,T,H]
        presents = [] if use_cache else None
        for i, layer in enumerate(self.layers):
            past = past_key_values[i] if past_key_values is not None else None
            x = layer(x, position_ids, token_embed, self.ring_bank,
                      past_kv=past, use_cache=use_cache, phase_seed=phase_seed)
            if use_cache:
                x, present = x
                presents.append(present)
        if self.hc_out_mix is not None:
            x = self.hc_out_mix(x)               # learned mix over streams (init == mean)
        # trunk refinements (no-op at init): HRM iterative gated refine, MoE expert mix
        rec = _viz.get_rec()
        if self.hrm is not None:
            hx = self.hrm(x)
            if rec is not None and rec.enabled:
                rec.log_trait("hrm", (hx - x)[0, -1].norm().item())
            x = hx
        if self.moe is not None:
            mx = self.moe(x)
            if rec is not None and rec.enabled:
                rec.log_trait("moe", mx[0, -1].norm().item())
            x = x + mx
        return self.norm_f(x), presents          # v2: RMSNorm trunk (no tanh)

    @torch.no_grad()
    def forward_drafts(self, idx, past_key_values=None, use_cache=False):
        """Speculative-decoding support (DeepSpec-style draft+verify, self-drafted):
        return (main_logits, mtp_logits_list[, presents]) over ALL T positions.
          * main_logits[:, j]        predicts the token at position j+1   (verifier)
          * mtp_logits_list[k][:, j] predicts the token at position j+2+k  (k=0..mtp-1)
        So one forward yields, at the last position, the genuine next token (main)
        plus `mtp_layers` drafted future tokens to be verified next round."""
        trunk, presents = self._compute_trunk(idx, past_key_values, use_cache)
        main_logits = self.head(trunk) * self.logit_scale
        mtp_logits = ([self.head(h(trunk)) * self.logit_scale for h in self.mtp_heads]
                      if self.mtp_heads is not None else [])
        return (main_logits, mtp_logits, presents) if use_cache else (main_logits, mtp_logits)

    def forward(self, idx, targets=None, past_key_values=None, use_cache=False):
        B, T = idx.shape
        trunk, presents = self._compute_trunk(idx, past_key_values, use_cache)
        logits = self.head(trunk) * self.logit_scale

        loss, aux = None, {}
        if targets is not None:
            flat = logits.view(-1, logits.size(-1))
            tflat = targets.reshape(-1)
            loss = F.cross_entropy(flat, tflat, ignore_index=-1)
            # z-loss (v2): penalise log^2 of the partition function -> no logit drift
            if self.cfg.zloss_coef > 0:
                valid = tflat != -1
                if valid.any():
                    log_z = torch.logsumexp(flat[valid].float(), dim=-1)
                    aux["zloss"] = self.cfg.zloss_coef * (log_z ** 2).mean()
            if self.moe is not None and self.moe.last_aux_loss is not None:
                aux["moe_aux"] = self.moe.last_aux_loss
            if self.mtp_heads is not None:
                mtp_total, n_active = 0.0, 0
                for k, head in enumerate(self.mtp_heads, start=1):
                    if T - k <= 0:
                        break
                    mlog = self.head(head(trunk[:, :T - k])) * self.logit_scale
                    mtp_total = mtp_total + F.cross_entropy(
                        mlog.reshape(-1, self.cfg.vocab_size),
                        targets[:, k:].reshape(-1), ignore_index=-1)
                    n_active += 1
                if n_active:
                    aux["mtp_loss"] = self.cfg.mtp_loss_weight * mtp_total / n_active
            if self.jepa is not None and T > 1:
                jepa_total, n_j = 0.0, 0
                for k in range(1, self.cfg.jepa_horizon + 1):
                    if T - k <= 0:
                        break
                    pred = self.jepa(trunk[:, :T - k], k)
                    tgt = trunk[:, k:].detach()      # JEPA stop-grad target
                    cos = F.cosine_similarity(pred.float(), tgt.float(), dim=-1)
                    jepa_total = jepa_total + (1.0 - cos).mean()
                    n_j += 1
                if n_j:
                    aux["jepa_loss"] = self.cfg.jepa_loss_weight * jepa_total / n_j
        # use_cache returns the per-layer (k, v) presents as a 4th item; without
        # it the signature is the original 3-tuple so train.py is unaffected.
        if use_cache:
            return logits, loss, aux, presents
        return logits, loss, aux

    @torch.no_grad()
    def generate(self, idx, n_new, temperature=1.0, top_k=None, use_cache=True):
        """Autoregressive sampling with the family's NaN/inf sanitisation: long
        accumulation can push a logit to +/-inf or NaN, whose softmax -> NaN ->
        multinomial fires a CUDA device-side assert. Clamp first, greedy-fallback.

        With use_cache=True (default) the attention KV cache is kept across steps
        so each new token costs one forward over a single position instead of a
        full recompute. The cache holds absolute positions, so once it fills past
        block_size we drop the cache and fall back to a windowed recompute (RoPE
        positions would otherwise exceed max_position_embeddings)."""
        self.eval()
        past = None
        for _ in range(n_new):
            if use_cache and past is not None and past[0][0].size(2) < self.cfg.block_size:
                step_in = idx[:, -1:]                       # only the newest token
                logits, _, _, past = self(step_in, past_key_values=past, use_cache=True)
            elif use_cache:
                cond = idx[:, -self.cfg.block_size:]        # prefill / re-prime cache
                logits, _, _, past = self(cond, use_cache=True)
            else:
                cond = idx[:, -self.cfg.block_size:]
                logits, _, _ = self(cond)
            lg = torch.nan_to_num(logits[:, -1, :].float(), nan=0.0,
                                  posinf=1e4, neginf=-1e4) / max(temperature, 1e-6)
            if top_k:
                k = min(top_k, lg.size(-1))
                v, _ = torch.topk(lg, k)
                lg[lg < v[:, [-1]]] = -float("inf")
            probs = torch.softmax(lg, dim=-1)
            if not torch.isfinite(probs).all() or float(probs.sum()) <= 0.0:
                nxt = torch.argmax(lg, dim=-1, keepdim=True)
            else:
                nxt = torch.multinomial(probs, 1)
            idx = torch.cat([idx, nxt], dim=1)
        return idx


if __name__ == "__main__":
    # exercise the base model and all trait modules at once
    cfg = QuazimotoConfig(use_hrm=True, use_moe=True, use_mtp=True, use_jepa=True)
    model = QuazimotoLM(cfg)
    x = torch.randint(0, cfg.vocab_size, (2, 64))
    logits, loss, aux = model(x, x)
    total = loss + sum(aux.values())
    print("forward ok:", tuple(logits.shape), "loss", round(loss.item(), 3),
          "| aux:", {k: round(float(v), 4) for k, v in aux.items()})
    total.backward()
    print("backward ok")