Blame for o1_native/hypercodec.py

"""Optical Hypercodec Language Layer, Kerr-Cavity O(1) Attention & Compound Carrier.

Solves all four small-SSM language & recall bottlenecks natively in fixed O(1) SRAM:
1. Universal Zero-UNK UTF-8 + 4,096-Crystal Bilingual/Code/Math Codec (O(1) bucketed
   longest-match encoder, 0.0% <UNK> on DE/EN/Code/Math).
2. Multi-Field O(1) Compound Carrier:
   - Exact ring-carried causal short-conv (`conv_ring`)
   - 4-plane GSSM-Selective log-complement rapidity stack (`depth_z`)
   - PS-Lifted Z2 Möbius-Lorentz polarity & momentum (`polarity_rho`)
   - 4-aperture Holographic Fourier phase waves (`holo_re`, `holo_im`)
   - Causal 1-step delayed Optical Kerr-Cavity Delta Attention (`cavity_prev_h`, `cavity_m`)
3. Non-Mutating Orthogonal Receiver Bank: parallel lenses read disjoint carrier
   subspaces (Syntax Shadow, Semantic Shadow, Focal Hierarchy, Z2 Polarity, Fourier
   Apertures, Kerr-Cavity Associative Readout, Epistemic Transmission Gate) without
   mutating the carrier (cryptographic SHA-256 invariant).
4. Protected Mathematical Invariant Embedding Subspace (`embed.weight[:, -16:]`)
   enabling direct tensor-level Organ Shells and zero-shot cross-lingual invariants.
"""

from __future__ import annotations

import hashlib
import math
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Sequence, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from .scan import (
    EPS,
    LOG_COMPLEMENT_CLAMP,
    MuPLinear,
    dispatch_affine_scan,
    holographic_aperture_scan,
    kerr_cavity_scan,
    moebius_rapidity_scan,
)


class ZeroUNKCodec:
    """Lossless UTF-8 byte + O(1)-indexed multi-byte crystal codec with 0% <UNK>.

    Token layout:
      0..255     : Raw UTF-8 bytes (guarantees every string on Earth is representable)
      256..263   : <BOS>, <EOS>, <PAD>, <SEP>, <FACT>, <SAY>, <SILENCE>, <QUERY>
      264..cap-1 : Lossless multi-byte vocabulary / morpheme / M-Rail crystal slots
    """

    BOS = 256
    EOS = 257
    PAD = 258
    SEP = 259
    FACT = 260
    SAY = 261
    SILENCE = 262
    QUERY = 263
    BASE_VOCAB = 264

    SPECIAL_MAP: Dict[str, int] = {
        "<BOS>": BOS,
        "<EOS>": EOS,
        "<PAD>": PAD,
        "<SEP>": SEP,
        "<FACT>": FACT,
        "<SAY>": SAY,
        "<SILENCE>": SILENCE,
        "<QUERY>": QUERY,
    }

    def __init__(self, vocab_capacity: int = 4096) -> None:
        if vocab_capacity < self.BASE_VOCAB:
            raise ValueError(f"vocab_capacity must be >= {self.BASE_VOCAB}")
        self.vocab_capacity = vocab_capacity
        self.id_to_bytes: Dict[int, bytes] = {i: bytes([i]) for i in range(256)}
        self.bytes_to_id: Dict[bytes, int] = {}
        self.next_crystal_id = self.BASE_VOCAB
        self._inv_special = {v: k for k, v in self.SPECIAL_MAP.items()}
        # 256 first-byte buckets sorted descending by byte length for fast O(1) greedy match
        self._first_byte_buckets: List[List[Tuple[bytes, int]]] = [[] for _ in range(256)]

    def _rebuild_buckets(self) -> None:
        buckets: List[List[Tuple[bytes, int]]] = [[] for _ in range(256)]
        for raw, cid in self.bytes_to_id.items():
            if len(raw) >= 2:
                buckets[raw[0]].append((raw, cid))
        for b_list in buckets:
            b_list.sort(key=lambda pair: len(pair[0]), reverse=True)
        self._first_byte_buckets = buckets

    def register_crystal(self, text: str) -> Optional[int]:
        """Crystallize a recurring multi-byte morpheme/word into a single token ID."""
        raw = text.encode("utf-8")
        if len(raw) <= 1:
            return raw[0] if len(raw) == 1 else None
        if raw in self.bytes_to_id:
            return self.bytes_to_id[raw]
        if self.next_crystal_id >= self.vocab_capacity:
            return None
        cid = self.next_crystal_id
        self.next_crystal_id += 1
        self.id_to_bytes[cid] = raw
        self.bytes_to_id[raw] = cid
        bucket = self._first_byte_buckets[raw[0]]
        bucket.append((raw, cid))
        bucket.sort(key=lambda pair: len(pair[0]), reverse=True)
        return cid

    def register_crystal_bytes(self, cid: int, raw: bytes) -> None:
        """Direct restoration helper used by snapshot loader."""
        self.id_to_bytes[cid] = raw
        if len(raw) >= 2:
            self.bytes_to_id[raw] = cid
            bucket = self._first_byte_buckets[raw[0]]
            bucket.append((raw, cid))
            bucket.sort(key=lambda pair: len(pair[0]), reverse=True)
        self.next_crystal_id = max(self.next_crystal_id, cid + 1)

    def encode(
        self,
        text: str,
        *,
        add_bos: bool = False,
        add_eos: bool = False,
        use_crystals: bool = True,
    ) -> List[int]:
        ids: List[int] = []
        if add_bos:
            ids.append(self.BOS)

        idx = 0
        n = len(text)
        has_crystals = use_crystals and bool(self.bytes_to_id)
        buckets = self._first_byte_buckets

        while idx < n:
            matched_special = None
            if text[idx] == "<":
                for marker, tok_id in self.SPECIAL_MAP.items():
                    if text.startswith(marker, idx):
                        matched_special = (marker, tok_id)
                        break
            if matched_special is not None:
                ids.append(matched_special[1])
                idx += len(matched_special[0])
                continue

            next_lt = text.find("<", idx + 1)
            chunk = text[idx:] if next_lt == -1 else text[idx:next_lt]
            raw = chunk.encode("utf-8")
            if not has_crystals:
                ids.extend(raw)
            else:
                b_idx = 0
                b_len = len(raw)
                while b_idx < b_len:
                    first_b = raw[b_idx]
                    candidates = buckets[first_b]
                    hit_id = None
                    hit_span = 1
                    rem_len = b_len - b_idx
                    for cand_bytes, cand_id in candidates:
                        c_len = len(cand_bytes)
                        if c_len <= rem_len and raw[b_idx : b_idx + c_len] == cand_bytes:
                            hit_id = cand_id
                            hit_span = c_len
                            break
                    if hit_id is not None:
                        ids.append(hit_id)
                        b_idx += hit_span
                    else:
                        ids.append(first_b)
                        b_idx += 1
            idx = n if next_lt == -1 else next_lt

        if add_eos:
            ids.append(self.EOS)
        return ids

    def decode(self, ids: Sequence[int], *, skip_special: bool = False) -> str:
        out_parts: List[str] = []
        byte_buf = bytearray()

        def flush_bytes() -> None:
            if byte_buf:
                out_parts.append(byte_buf.decode("utf-8", errors="replace"))
                byte_buf.clear()

        for tok in ids:
            t = int(tok)
            if t in self.id_to_bytes:
                byte_buf.extend(self.id_to_bytes[t])
            elif t in self._inv_special:
                flush_bytes()
                if not skip_special:
                    out_parts.append(self._inv_special[t])
            elif 0 <= t < 256:
                byte_buf.append(t)
        flush_bytes()
        return "".join(out_parts)

    def compression_stats(self, text: str) -> Dict[str, float]:
        raw_len = len(text.encode("utf-8"))
        ids = self.encode(text, use_crystals=True)
        cry_len = len(ids)
        unk_cnt = sum(
            1 for tid in ids if tid not in self.id_to_bytes and tid not in self._inv_special
        )
        return {
            "utf8_bytes": float(raw_len),
            "crystal_tokens": float(cry_len),
            "bytes_per_token": round(raw_len / float(max(cry_len, 1)), 4),
            "active_crystals": float(len(self.bytes_to_id)),
            "unk_count": float(unk_cnt),
        }


@dataclass(frozen=True)
class CarrierConfig:
    vocab_size: int = 4096
    d_model: int = 192
    n_layers: int = 2
    n_depths: int = 4          # byte/syllable, word, clause, discourse
    d_depth: int = 16
    n_polarities: int = 4      # PS-Lifted Z2 polarity channels
    d_polarity: int = 8
    n_apertures: int = 4       # Disjoint holographic Fourier phase windows
    d_aperture: int = 16
    n_cavity_heads: int = 4    # Optical Kerr-Cavity O(1) Attention heads
    cavity_angles: int = 8     # P = 4 * cavity_angles = 32 orthogonal Fejér modes/head
    d_cavity_val: int = 24     # Value channels per cavity head
    use_kerr_resonator: bool = True
    invariant_dim: int = 16    # Protected tail coordinates in embed.weight
    conv_kernel: int = 4       # Causal short-conv kernel (exact ring-carried)
    ffn_mult: int = 2
    focal_depths: Tuple[float, ...] = (-1.0, -0.33, 0.33, 1.0)
    depth_gamma_init: Tuple[float, ...] = (0.30, 0.68, 0.90, 0.992)

    @classmethod
    def ane_m4_p64(
        cls,
        vocab_size: int = 4096,
        d_model: int = 192,
        n_layers: int = 2,
    ) -> "CarrierConfig":
        """Hardware-quantum M4 ANE configuration (`P = 4 * 16 = 64`, `d_cavity_val = 64`)."""
        return cls(
            vocab_size=int(vocab_size),
            d_model=int(d_model),
            n_layers=int(n_layers),
            cavity_angles=16,
            d_cavity_val=64,
            conv_kernel=8,
        )


@dataclass
class LayerCarrierState:
    """Exact O(1) recurrent carrier state of a single HypercodecLayer."""

    conv_ring: torch.Tensor       # (B, d_model, conv_kernel - 1)
    depth_z: torch.Tensor         # (B, n_depths, d_depth) in (-inf, 0]
    polarity_rho: torch.Tensor    # (B, n_polarities, d_polarity) in (-inf, +inf)
    holo_re: torch.Tensor         # (B, n_apertures, d_aperture)
    holo_im: torch.Tensor         # (B, n_apertures, d_aperture)
    cavity_prev_h: torch.Tensor   # (B, d_model) exact 1-step causal delay register
    cavity_m: torch.Tensor        # (B, n_cavity_heads, 4 * cavity_angles, d_cavity_val)

    def detach(self) -> "LayerCarrierState":
        return LayerCarrierState(
            conv_ring=self.conv_ring.detach(),
            depth_z=self.depth_z.detach(),
            polarity_rho=self.polarity_rho.detach(),
            holo_re=self.holo_re.detach(),
            holo_im=self.holo_im.detach(),
            cavity_prev_h=self.cavity_prev_h.detach(),
            cavity_m=self.cavity_m.detach(),
        )

    def clone(self) -> "LayerCarrierState":
        return LayerCarrierState(
            conv_ring=self.conv_ring.clone(),
            depth_z=self.depth_z.clone(),
            polarity_rho=self.polarity_rho.clone(),
            holo_re=self.holo_re.clone(),
            holo_im=self.holo_im.clone(),
            cavity_prev_h=self.cavity_prev_h.clone(),
            cavity_m=self.cavity_m.clone(),
        )

    def tensors(self) -> Tuple[torch.Tensor, ...]:
        return (
            self.conv_ring,
            self.depth_z,
            self.polarity_rho,
            self.holo_re,
            self.holo_im,
            self.cavity_prev_h,
            self.cavity_m,
        )


@dataclass
class CompoundCarrierState:
    """Complete O(1) compound carrier across all layers + last surface readout slice."""

    layers: List[LayerCarrierState]
    surface_field: Optional[Dict[str, torch.Tensor]] = None
    step_count: int = 0

    def detach(self) -> "CompoundCarrierState":
        surf = None
        if self.surface_field is not None:
            surf = {k: v.detach() for k, v in self.surface_field.items()}
        return CompoundCarrierState(
            layers=[layer.detach() for layer in self.layers],
            surface_field=surf,
            step_count=self.step_count,
        )

    def clone(self) -> "CompoundCarrierState":
        surf = None
        if self.surface_field is not None:
            surf = {k: v.clone() for k, v in self.surface_field.items()}
        return CompoundCarrierState(
            layers=[layer.clone() for layer in self.layers],
            surface_field=surf,
            step_count=self.step_count,
        )

    def nbytes(self) -> int:
        total = 0
        for layer in self.layers:
            for t in layer.tensors():
                total += t.numel() * t.element_size()
        return total

    def sha256(self) -> str:
        digest = hashlib.sha256()
        digest.update(f"step:{self.step_count}".encode("ascii"))
        for idx, layer in enumerate(self.layers):
            digest.update(f"L{idx}".encode("ascii"))
            for t in layer.tensors():
                arr = np.ascontiguousarray(t.detach().cpu().numpy())
                digest.update(str(arr.dtype).encode("ascii"))
                digest.update(repr(arr.shape).encode("ascii"))
                digest.update(arr.view(np.uint8).tobytes())
        if self.surface_field is not None:
            for key in sorted(self.surface_field.keys()):
                digest.update(key.encode("utf-8"))
                arr = np.ascontiguousarray(self.surface_field[key].detach().cpu().numpy())
                digest.update(arr.view(np.uint8).tobytes())
        return digest.hexdigest()

    def pack_flat_f32(self) -> np.ndarray:
        """Pack recurrent O(1) state into one contiguous float32 array for NPU SRAM."""
        chunks = []
        for layer in self.layers:
            for t in layer.tensors():
                chunks.append(t.detach().cpu().to(torch.float32).reshape(-1).numpy())
        return np.concatenate(chunks, axis=0) if chunks else np.zeros((0,), dtype=np.float32)


class HypercodecLayer(nn.Module):
    """Single NoPE Multi-Field GSSM + Kerr-Cavity O(1) Attention + Hypercodec Layer."""

    def __init__(self, cfg: CarrierConfig) -> None:
        super().__init__()
        self.cfg = cfg
        D = cfg.d_model
        K = cfg.conv_kernel

        # 1. Exact ring-carried depthwise causal short-conv
        self.conv_dw = nn.Conv1d(D, D, kernel_size=K, groups=D, bias=True)
        nn.init.xavier_uniform_(self.conv_dw.weight, gain=0.4)
        nn.init.zeros_(self.conv_dw.bias)

        # 2. Multi-depth bounded selective rapidity projections (muP scaled)
        depth_dim = cfg.n_depths * cfg.d_depth
        self.W_depth_v = MuPLinear(D, depth_dim, base_d=D, init_gain=0.6)
        self.W_depth_gate = MuPLinear(D, depth_dim, base_d=D, init_gain=0.6)
        self.W_depth_gamma = MuPLinear(D, depth_dim, base_d=D, init_gain=0.15, bias=True)
        self.W_depth_alpha = MuPLinear(D, depth_dim, base_d=D, init_gain=0.15, bias=True)

        with torch.no_grad():
            bias_view = self.W_depth_gamma.linear.bias.view(cfg.n_depths, cfg.d_depth)
            for idx in range(cfg.n_depths):
                g_init = cfg.depth_gamma_init[idx % len(cfg.depth_gamma_init)]
                g_clamped = min(max(g_init, 0.01), 0.999)
                logit = math.log(g_clamped / (1.0 - g_clamped))
                bias_view[idx].fill_(logit)

        # 3. PS-Lifted Z2 Möbius-Lorentz polarity projections (high clause persistence + sparse write)
        pol_dim = cfg.n_polarities * cfg.d_polarity
        self.W_pol_v = MuPLinear(D, pol_dim, base_d=D, init_gain=0.6)
        self.W_pol_gamma = MuPLinear(D, pol_dim, base_d=D, init_gain=0.15, bias=True)
        self.W_pol_beta = MuPLinear(D, pol_dim, base_d=D, init_gain=0.4, bias=True)
        with torch.no_grad():
            self.W_pol_gamma.linear.bias.fill_(math.log(0.96 / 0.04))
            self.W_pol_beta.linear.bias.fill_(-2.0)

        # 4. Multi-aperture holographic complex phase plate (separate Q/K angles + sparse write)
        holo_dim = cfg.n_apertures * cfg.d_aperture
        self.W_holo_u = MuPLinear(D, holo_dim, base_d=D, init_gain=0.6)
        self.W_holo_gate = MuPLinear(D, holo_dim, base_d=D, init_gain=0.4, bias=True)
        self.W_holo_gamma = MuPLinear(D, holo_dim, base_d=D, init_gain=0.15, bias=True)
        self.W_holo_k_phase = MuPLinear(D, holo_dim, base_d=D, init_gain=0.25)
        self.W_holo_q_phase = MuPLinear(D, holo_dim, base_d=D, init_gain=0.25)
        with torch.no_grad():
            self.W_holo_gamma.linear.bias.fill_(math.log(0.96 / 0.04))
            self.W_holo_gate.linear.bias.fill_(-2.0)

        # 5. Native O(1) Optical Kerr-Cavity Delta Attention
        #    Shared orthogonal anchor guarantees step-0 induction alignment:
        #    theta_k(t) reads h_{t-1} (key before value), theta_q(T) reads h_T (query key)!
        cav_ang_dim = cfg.n_cavity_heads * cfg.cavity_angles
        cav_val_dim = cfg.n_cavity_heads * cfg.d_cavity_val
        self.W_cav_anchor = nn.Linear(D, cav_ang_dim, bias=False)
        nn.init.orthogonal_(self.W_cav_anchor.weight, gain=1.1)
        self.W_cav_k_ctx = MuPLinear(D, cav_ang_dim, base_d=D, init_gain=0.06)
        self.W_cav_q_ctx = MuPLinear(D, cav_ang_dim, base_d=D, init_gain=0.06)
        self.W_cav_val = MuPLinear(2 * D, cav_val_dim, base_d=2 * D, init_gain=0.7)
        self.W_cav_beta = MuPLinear(2 * D, cfg.n_cavity_heads, base_d=2 * D, init_gain=0.4, bias=True)
        self.W_cav_erase = MuPLinear(2 * D, cfg.n_cavity_heads, base_d=2 * D, init_gain=0.3, bias=True)
        self.W_cav_decay = MuPLinear(2 * D, cfg.n_cavity_heads, base_d=2 * D, init_gain=0.2, bias=True)
        with torch.no_grad():
            self.W_cav_beta.linear.bias.fill_(-1.2)   # Selective write
            self.W_cav_erase.linear.bias.fill_(1.8)   # Strong Delta erase-before-write
            self.W_cav_decay.linear.bias.fill_(-3.0)  # Near-lossless retention over gaps

        if cfg.use_kerr_resonator:
            self.resonator_lam = nn.Parameter(torch.full((1, cfg.n_cavity_heads, 1), 0.02))
        else:
            self.register_parameter("resonator_lam", None)

        total_carrier_read = depth_dim + 2 * pol_dim + 2 * holo_dim + cav_val_dim
        self.norm_pre = nn.RMSNorm(D)
        self.W_merge = MuPLinear(total_carrier_read, D, base_d=D, init_gain=0.6)
        self.norm_ffn = nn.RMSNorm(D)
        ffn_hidden = D * cfg.ffn_mult
        self.ffn_gate = nn.Linear(D, ffn_hidden, bias=False)
        self.ffn_up = nn.Linear(D, ffn_hidden, bias=False)
        self.ffn_down = nn.Linear(ffn_hidden, D, bias=False)

    def init_state(
        self,
        batch_size: int,
        device: torch.device,
        dtype: torch.dtype = torch.float32,
    ) -> LayerCarrierState:
        cfg = self.cfg
        return LayerCarrierState(
            conv_ring=torch.zeros(
                batch_size, cfg.d_model, max(cfg.conv_kernel - 1, 1), device=device, dtype=dtype
            ),
            depth_z=torch.zeros(
                batch_size, cfg.n_depths, cfg.d_depth, device=device, dtype=dtype
            ),
            polarity_rho=torch.zeros(
                batch_size, cfg.n_polarities, cfg.d_polarity, device=device, dtype=dtype
            ),
            holo_re=torch.zeros(
                batch_size, cfg.n_apertures, cfg.d_aperture, device=device, dtype=dtype
            ),
            holo_im=torch.zeros(
                batch_size, cfg.n_apertures, cfg.d_aperture, device=device, dtype=dtype
            ),
            cavity_prev_h=torch.zeros(batch_size, cfg.d_model, device=device, dtype=dtype),
            cavity_m=torch.zeros(
                batch_size,
                cfg.n_cavity_heads,
                4 * cfg.cavity_angles,
                cfg.d_cavity_val,
                device=device,
                dtype=dtype,
            ),
        )

    def forward(
        self,
        x: torch.Tensor,
        state: Optional[LayerCarrierState] = None,
        *,
        scan_mode: str = "auto",
    ) -> Tuple[torch.Tensor, LayerCarrierState, Dict[str, torch.Tensor]]:
        B, T, D = x.shape
        cfg = self.cfg
        if state is None:
            state = self.init_state(B, x.device, x.dtype)

        h = self.norm_pre(x)

        # 1. Exact ring-carried causal short convolution
        K = cfg.conv_kernel
        h_t = h.transpose(1, 2)  # (B, D, T)
        if K > 1:
            padded = torch.cat([state.conv_ring, h_t], dim=2)
            new_conv_ring = padded[:, :, -(K - 1) :]
            u_local = F.silu(self.conv_dw(padded)).transpose(1, 2)
        else:
            new_conv_ring = state.conv_ring
            u_local = F.silu(self.conv_dw(h_t)).transpose(1, 2)

        # 2. Multi-Depth GSSM-Selective log-complement rapidity scan
        v_d = torch.tanh(self.W_depth_v(u_local))
        g_d = torch.sigmoid(self.W_depth_gate(u_local))
        gamma_d = torch.sigmoid(self.W_depth_gamma(u_local)).view(
            B, T, cfg.n_depths, cfg.d_depth
        )
        alpha_d = torch.sigmoid(self.W_depth_alpha(u_local)).view(
            B, T, cfg.n_depths, cfg.d_depth
        )
        v_gated = (v_d * g_d).view(B, T, cfg.n_depths, cfg.d_depth)
        w_sq = torch.clamp(v_gated * v_gated, max=LOG_COMPLEMENT_CLAMP)
        drive_d = alpha_d * torch.log(1.0 - w_sq + EPS)
        z_seq, z_fin = dispatch_affine_scan(drive_d, gamma_d, state.depth_z, mode=scan_mode)
        depth_s = torch.sqrt(torch.clamp(1.0 - torch.exp(z_seq), min=0.0) + EPS)

        # 3. PS-Lifted Z2 Möbius-Lorentz polarity & momentum scan
        v_p = torch.tanh(self.W_pol_v(u_local)).view(
            B, T, cfg.n_polarities, cfg.d_polarity
        )
        gamma_p = torch.sigmoid(self.W_pol_gamma(u_local)).view(
            B, T, cfg.n_polarities, cfg.d_polarity
        )
        beta_p = torch.sigmoid(self.W_pol_beta(u_local)).view(
            B, T, cfg.n_polarities, cfg.d_polarity
        )
        pol_seq, _rho_seq, rho_fin = moebius_rapidity_scan(
            v_p, gamma_p, beta_p, state.polarity_rho, scan_mode=scan_mode
        )
        pol_mag = depth_s.mean(dim=-1, keepdim=True).mean(dim=-2, keepdim=True)
        z2_plus = 0.5 * (1.0 + pol_seq) * (0.5 + 0.5 * pol_mag)
        z2_minus = 0.5 * (1.0 - pol_seq) * (0.5 + 0.5 * pol_mag)
        z2_physical = z2_plus + z2_minus
        z2_momentum = z2_plus - z2_minus

        # 4. Multi-Aperture Holographic Fourier Phase Scan
        u_holo_raw = torch.tanh(self.W_holo_u(u_local)) * torch.sigmoid(
            self.W_holo_gate(u_local)
        )
        u_holo = u_holo_raw.view(B, T, cfg.n_apertures, cfg.d_aperture)
        gamma_h = torch.sigmoid(self.W_holo_gamma(u_local)).view(
            B, T, cfg.n_apertures, cfg.d_aperture
        )
        phi_k = math.pi * torch.tanh(self.W_holo_k_phase(u_local)).view(
            B, T, cfg.n_apertures, cfg.d_aperture
        )
        phi_q = math.pi * torch.tanh(self.W_holo_q_phase(u_local)).view(
            B, T, cfg.n_apertures, cfg.d_aperture
        )
        holo_re_read, holo_im_read, holo_re_fin, holo_im_fin = holographic_aperture_scan(
            u_holo, gamma_h, phi_k, phi_q, state.holo_re, state.holo_im, scan_mode=scan_mode
        )
        rms_re = torch.sqrt(torch.mean(holo_re_read * holo_re_read, dim=-1, keepdim=True) + 1e-5)
        rms_im = torch.sqrt(torch.mean(holo_im_read * holo_im_read, dim=-1, keepdim=True) + 1e-5)
        holo_re_norm = holo_re_read / rms_re
        holo_im_norm = holo_im_read / rms_im

        # 5. Native O(1) Optical Kerr-Cavity Delta Attention (exact carried 1-step shift)
        if T > 0:
            h_prev_seq = torch.cat([state.cavity_prev_h.unsqueeze(1), h[:, :-1, :]], dim=1)
            new_cavity_prev_h = h[:, -1, :]
        else:
            h_prev_seq = h
            new_cavity_prev_h = state.cavity_prev_h

        H_c = cfg.n_cavity_heads
        M_ang = cfg.cavity_angles
        D_v = cfg.d_cavity_val
        pair_ctx = torch.cat([h_prev_seq, h], dim=-1)

        theta_k = math.pi * torch.tanh(
            self.W_cav_anchor(h_prev_seq) + self.W_cav_k_ctx(u_local)
        ).view(B, T, H_c, M_ang)
        theta_q = math.pi * torch.tanh(
            self.W_cav_anchor(h) + self.W_cav_q_ctx(u_local)
        ).view(B, T, H_c, M_ang)

        cav_vals = torch.tanh(self.W_cav_val(pair_ctx)).view(B, T, H_c, D_v)
        cav_beta = torch.sigmoid(self.W_cav_beta(pair_ctx)).view(B, T, H_c, 1)
        cav_erase = torch.sigmoid(self.W_cav_erase(pair_ctx)).view(B, T, H_c, 1)
        cav_decay = (0.15 * torch.sigmoid(self.W_cav_decay(pair_ctx))).view(B, T, H_c, 1)

        cav_read_seq, cav_m_fin = kerr_cavity_scan(
            theta_k,
            theta_q,
            cav_vals,
            cav_beta,
            cav_erase,
            cav_decay,
            state.cavity_m,
            resonator_lam=self.resonator_lam,
        )

        merged_features = torch.cat(
            [
                depth_s.reshape(B, T, -1),
                z2_physical.reshape(B, T, -1),
                z2_momentum.reshape(B, T, -1),
                holo_re_norm.reshape(B, T, -1),
                holo_im_norm.reshape(B, T, -1),
                cav_read_seq.reshape(B, T, -1),
            ],
            dim=-1,
        )
        y = x + self.W_merge(merged_features)
        h_ffn = self.norm_ffn(y)
        f_act = F.silu(self.ffn_gate(h_ffn)) * self.ffn_up(h_ffn)
        y = y + self.ffn_down(f_act)

        next_state = LayerCarrierState(
            conv_ring=new_conv_ring,
            depth_z=z_fin,
            polarity_rho=rho_fin,
            holo_re=holo_re_fin,
            holo_im=holo_im_fin,
            cavity_prev_h=new_cavity_prev_h,
            cavity_m=cav_m_fin,
        )
        carrier_slice = {
            "depth_s": depth_s[:, -1],
            "z2_physical": z2_physical[:, -1],
            "z2_momentum": z2_momentum[:, -1],
            "holo_re": holo_re_norm[:, -1],
            "holo_im": holo_im_norm[:, -1],
            "cavity_read": cav_read_seq[:, -1],
            "seq_depth_s": depth_s,
            "seq_z2_momentum": z2_momentum,
            "seq_holo_re": holo_re_norm,
            "seq_cavity_read": cav_read_seq,
        }
        return y, next_state, carrier_slice


class OrthogonalReceiverBank(nn.Module):
    """Non-mutating parallel Optical Hypercodec receivers over the Compound Carrier."""

    def __init__(self, cfg: CarrierConfig) -> None:
        super().__init__()
        self.cfg = cfg
        D = cfg.d_model
        inv_d = cfg.invariant_dim if D >= 64 else 0
        usable_d = D - inv_d
        half_d = usable_d // 2
        self.half_d = half_d
        self.usable_d = usable_d
        self.inv_d = inv_d

        self.norm_final = nn.RMSNorm(D)

        # Disjoint orthogonal shadow projections (XY syntax vs XZ relation; excludes invariant tail)
        self.W_shadow_xy = nn.Linear(half_d, cfg.vocab_size, bias=False)
        self.W_shadow_xz = nn.Linear(usable_d - half_d, cfg.vocab_size, bias=False)

        # Focal-depth hierarchy lens -> vocabulary contribution
        self.W_focus_vocab = nn.Linear(cfg.d_depth, cfg.vocab_size, bias=False)

        # Z2 polarization momentum lens -> directional / polarity contribution
        pol_flat = cfg.n_polarities * cfg.d_polarity
        self.W_pol_vocab = nn.Linear(pol_flat, cfg.vocab_size, bias=False)

        # Disjoint Fourier aperture lenses -> wave phase contribution
        holo_flat = cfg.n_apertures * cfg.d_aperture
        self.W_fourier_vocab = nn.Linear(holo_flat, cfg.vocab_size, bias=False)

        # Dedicated Optical Kerr-Cavity associative recall lens
        cav_flat = cfg.n_cavity_heads * cfg.d_cavity_val
        self.W_cavity_vocab = nn.Linear(cav_flat, cfg.vocab_size, bias=False)

        # Epistemic transmission / Foss-Gate (measures carrier coherence -> [0, 1])
        self.W_transmission = nn.Linear(D + pol_flat, 1, bias=True)
        nn.init.constant_(self.W_transmission.bias, 1.5)

        for head in (self.W_focus_vocab, self.W_pol_vocab, self.W_fourier_vocab):
            nn.init.xavier_uniform_(head.weight, gain=0.25)
        nn.init.xavier_uniform_(self.W_cavity_vocab.weight, gain=0.45)

    def render_focus_plane(
        self,
        depth_stack: torch.Tensor,
        focus_depth: float = 0.0,
        aperture_sharpness: float = 6.0,
    ) -> torch.Tensor:
        """Continuous optical focus sweep across the depth planes."""
        if len(self.cfg.focal_depths) == self.cfg.n_depths:
            depths = torch.tensor(
                self.cfg.focal_depths, device=depth_stack.device, dtype=depth_stack.dtype
            )
        else:
            depths = torch.linspace(
                -1.0, 1.0, self.cfg.n_depths, device=depth_stack.device, dtype=depth_stack.dtype
            )
        dist = torch.abs(depths - float(focus_depth))
        weights = 1.0 / (1.0 + aperture_sharpness * dist * dist)
        weights = weights / torch.clamp(weights.sum(), min=EPS)
        view_shape = (1,) * (depth_stack.ndim - 2) + (self.cfg.n_depths, 1)
        return (depth_stack * weights.view(view_shape)).sum(dim=-2)

    def inspect_all_lenses(
        self,
        carrier: CompoundCarrierState,
        *,
        focus_query: float = 0.0,
    ) -> Dict[str, Any]:
        """Parallel non-mutating readout of all optical receiver lenses + SHA proof."""
        hash_before = carrier.sha256()
        if carrier.surface_field is None:
            raise ValueError("Carrier has no surface_field; run at least 1 step first.")

        sf = carrier.surface_field
        h_last = sf["h_norm"]
        depth_s = sf["depth_s"]
        z2_mom = sf["z2_momentum"]
        holo_re = sf["holo_re"]
        holo_im = sf["holo_im"]
        cav_read = sf["cavity_read"]

        shadow_xy = h_last[:, : self.half_d].detach().clone()
        shadow_xz = h_last[:, self.half_d : self.usable_d].detach().clone()

        focus_slices = {
            f"focus.plane_{i}": depth_s[:, i, :].detach().clone()
            for i in range(self.cfg.n_depths)
        }
        focus_queried = self.render_focus_plane(depth_s, focus_depth=focus_query).detach().clone()

        fourier_windows = {
            f"fourier.window_{w}": torch.sqrt(
                holo_re[:, w, :] ** 2 + holo_im[:, w, :] ** 2 + EPS
            ).detach().clone()
            for w in range(self.cfg.n_apertures)
        }
        cavity_heads = {
            f"cavity.head_{h_idx}": cav_read[:, h_idx, :].detach().clone()
            for h_idx in range(self.cfg.n_cavity_heads)
        }

        pol_flat = z2_mom.reshape(h_last.shape[0], -1)
        transmission = torch.sigmoid(
            self.W_transmission(torch.cat([h_last, pol_flat], dim=-1))
        ).detach().clone()

        hash_after = carrier.sha256()
        return {
            "carrier_hash_before": hash_before,
            "carrier_hash_after": hash_after,
            "carrier_immutable": hash_before == hash_after,
            "shadow.xy_syntax": shadow_xy,
            "shadow.xz_relation": shadow_xz,
            "focus.queried": focus_queried,
            **focus_slices,
            "polarization.z2_momentum": z2_mom.detach().clone(),
            **fourier_windows,
            **cavity_heads,
            "surface.transmission_gate": transmission,
        }

    def forward(
        self,
        h_seq: torch.Tensor,
        seq_depth_s: torch.Tensor,
        seq_z2_momentum: torch.Tensor,
        seq_holo_re: torch.Tensor,
        seq_cavity_read: torch.Tensor,
        *,
        focus_depth: float = 0.0,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Factorized additive optical lens synthesis into vocabulary logits."""
        B, T, _ = h_seq.shape
        h_norm = self.norm_final(h_seq)
        xy = h_norm[:, :, : self.half_d]
        xz = h_norm[:, :, self.half_d : self.usable_d]

        focused = self.render_focus_plane(seq_depth_s, focus_depth=focus_depth)
        pol_flat = seq_z2_momentum.reshape(B, T, -1)
        holo_flat = seq_holo_re.reshape(B, T, -1)
        cav_flat = seq_cavity_read.reshape(B, T, -1)

        logits = (
            self.W_shadow_xy(xy)
            + self.W_shadow_xz(xz)
            + self.W_focus_vocab(focused)
            + self.W_pol_vocab(pol_flat)
            + self.W_fourier_vocab(holo_flat)
            + self.W_cavity_vocab(cav_flat)
        )
        transmission = torch.sigmoid(
            self.W_transmission(torch.cat([h_norm, pol_flat], dim=-1))
        )
        return logits, transmission, h_norm


class OrthogonalAttributeProbe(nn.Module):
    """Dedicated disjoint-channel Hypercodec Lenses for compositional binding (P91b proof)."""

    def __init__(self, cfg: CarrierConfig, n_dirs: int, n_mags: int) -> None:
        super().__init__()
        pol_dim = cfg.n_polarities * cfg.d_polarity
        ap_dim = cfg.d_aperture
        self.dir_lens = nn.Linear(pol_dim, n_dirs)
        self.mag_lens = nn.Linear(ap_dim, n_mags)

    def forward(self, carrier: CompoundCarrierState) -> Tuple[torch.Tensor, torch.Tensor]:
        if not carrier.layers:
            raise ValueError("Carrier layers missing")
        primary = carrier.layers[0]
        pol = torch.tanh(primary.polarity_rho).reshape(primary.polarity_rho.shape[0], -1)
        ap0 = primary.holo_re[:, 0, :]
        return self.dir_lens(pol), self.mag_lens(ap0)


class O1HypercodecModel(nn.Module):
    """Complete Native NoPE O1-GSSM + Kerr-Cavity + Optical Hypercodec Model."""

    def __init__(self, cfg: Optional[CarrierConfig] = None) -> None:
        super().__init__()
        self.cfg = cfg or CarrierConfig()
        self.embed = nn.Embedding(self.cfg.vocab_size, self.cfg.d_model)
        nn.init.normal_(self.embed.weight, mean=0.0, std=0.04)
        self._invariant_grad_hook_registered = False
        self.layers = nn.ModuleList(
            [HypercodecLayer(self.cfg) for _ in range(self.cfg.n_layers)]
        )
        self.receivers = OrthogonalReceiverBank(self.cfg)

    def lock_invariant_embedding_slice(self) -> None:
        """Freeze gradient drift on the protected mathematical invariant coordinates `[:, -16:]`."""
        inv_d = self.cfg.invariant_dim if self.cfg.d_model >= 64 else 0
        if inv_d <= 0 or self._invariant_grad_hook_registered:
            return

        def _mask_inv_grad(grad: torch.Tensor) -> torch.Tensor:
            if grad is None:
                return grad
            masked = grad.clone()
            masked[:, -inv_d:] = 0.0
            return masked

        self.embed.weight.register_hook(_mask_inv_grad)
        self._invariant_grad_hook_registered = True

    def init_carrier(
        self,
        batch_size: int = 1,
        device: Optional[torch.device] = None,
        dtype: torch.dtype = torch.float32,
    ) -> CompoundCarrierState:
        dev = device or self.embed.weight.device
        return CompoundCarrierState(
            layers=[layer.init_state(batch_size, dev, dtype) for layer in self.layers],
            surface_field=None,
            step_count=0,
        )

    def parameter_count(self) -> int:
        return sum(p.numel() for p in self.parameters())

    def forward(
        self,
        input_ids: torch.Tensor,
        carrier_in: Optional[CompoundCarrierState] = None,
        *,
        focus_depth: float = 0.0,
        scan_mode: str = "auto",
    ) -> Dict[str, Any]:
        B, T = input_ids.shape
        if carrier_in is None:
            carrier_in = self.init_carrier(B, input_ids.device, self.embed.weight.dtype)

        h = self.embed(input_ids)
        next_layers: List[LayerCarrierState] = []
        last_slice: Dict[str, torch.Tensor] = {}

        for layer, l_state in zip(self.layers, carrier_in.layers):
            h, l_next, c_slice = layer(h, l_state, scan_mode=scan_mode)
            next_layers.append(l_next)
            last_slice = c_slice

        logits, transmission, h_norm = self.receivers(
            h,
            last_slice["seq_depth_s"],
            last_slice["seq_z2_momentum"],
            last_slice["seq_holo_re"],
            last_slice["seq_cavity_read"],
            focus_depth=focus_depth,
        )

        surface = {
            "h_norm": h_norm[:, -1],
            "depth_s": last_slice["depth_s"],
            "z2_physical": last_slice["z2_physical"],
            "z2_momentum": last_slice["z2_momentum"],
            "holo_re": last_slice["holo_re"],
            "holo_im": last_slice["holo_im"],
            "cavity_read": last_slice["cavity_read"],
            "transmission": transmission[:, -1],
        }
        carrier_out = CompoundCarrierState(
            layers=next_layers,
            surface_field=surface,
            step_count=carrier_in.step_count + T,
        )
        return {
            "logits": logits,
            "transmission": transmission,
            "carrier": carrier_out,
        }