"""aleph_diffusion_core.py — runner-2 diffusion aleph-adapter core (exp001+).

The certified components ported to diffusion trunks (repos/geolip-aleph-diffusion.md):
  AlephAddress    — verbatim port of tools/aleph_routed_attention.py:51 (closed-form
                    sinh/cosh read over 2K oriented half-axes; canon/aleph_core.md).
  RelayPatch2D    — the RelayPatchwork port (tools/qwen_exp002_refine.py:96):
                    per-block residual relay on token-space hidden states [B,T,d].
                    Zero-init out weight AND bias (exp006 bias-leak law), gate -3.0,
                    enabled=False returns x via code-path skip (toggle law bit-exact).
  AlephCondAdapter— conditioning-stream address injection (geolip-sdxl-aleph
                    addr_adapter precedent): frozen [32,128] address -> 32 cond
                    positions, everything zero-init => exact zeros at init.
  Sign-code gauge — register-probe separation, DIAGONAL EXCLUDED (documented
                    discontinuity vs the qwen line; canon/register_probe_gauge.md).
  Parity utils    — P-TOGGLE / P-INIT / P-FIRE / P-SAVE gate helpers.

Riders: pure Adam wd=0 downstream; no comparative selectors in any gradient path
(argmax appears ONLY in read-only gauges); fp32 for paired gauges; no GAP on
feature paths (pooling appears ONLY in read-only gauges).

CPU cert:  python pod2/aleph_diffusion_core.py   (shapes/parity only — never
CPU-train for accuracy; house law).
"""
from __future__ import annotations

import math

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

PHI = math.sqrt(2.0)
PSI = 1.533751168755204288118041


def super_fibonacci_s3(n: int, dtype=torch.float32) -> torch.Tensor:
    """Near-uniform unit quaternions (Alexa CVPR'22), constants per canon."""
    i = torch.arange(n, dtype=torch.float64)
    s = (i + 0.5) / n
    r = torch.sqrt(s)
    R = torch.sqrt(1.0 - s)
    alpha = 2.0 * math.pi * i / PHI
    beta = 2.0 * math.pi * i / PSI
    q = torch.stack([r * torch.sin(alpha), r * torch.cos(alpha),
                     R * torch.sin(beta), R * torch.cos(beta)], dim=-1)
    return F.normalize(q, dim=-1).to(dtype)


class AlephAddress(nn.Module):
    """The aleph addresser over 2K oriented half-axes [+A; -A].
    m_hat(x): closed-form soft read Sum_k sinh(u_k)A_k / Sum_k cosh(u_k), stable
    via max-|u| factor-out; 2K never materialized. Codebook trains ONLY through
    whatever consumes the address (no commit/EMA/VQ)."""

    def __init__(self, K: int = 64, D: int = 4, tau: float = 0.1,
                 init: str = "random", codebook: torch.Tensor | None = None,
                 learnable: bool = True):
        super().__init__()
        self.K, self.D, self.tau = K, D, tau
        if codebook is not None:
            A = F.normalize(codebook.detach().clone().float(), dim=-1)
            assert A.shape == (K, D), f"custom codebook must be ({K},{D})"
        elif init == "fibonacci":
            assert D == 4, "fibonacci init is defined on S^3 (D=4) only"
            A = super_fibonacci_s3(K)
        elif init == "random":
            A = F.normalize(torch.randn(K, D), dim=-1)
        else:
            raise ValueError(f"unknown init '{init}'")
        if learnable:
            self.codebook = nn.Parameter(A)
        else:
            self.register_buffer("codebook", A)
        self.register_buffer("home", A.detach().clone())

    def _u(self, x: torch.Tensor) -> torch.Tensor:
        A = F.normalize(self.codebook, dim=-1)
        return (F.normalize(x, dim=-1) @ A.transpose(-1, -2)) / self.tau

    def m_hat(self, x: torch.Tensor) -> torch.Tensor:
        """Closed-form soft read (decoders read THIS, never the raw codebook)."""
        u = self._u(x)
        m = u.abs().amax(dim=-1, keepdim=True)
        ep, en = torch.exp(u - m), torch.exp(-u - m)
        num = ep - en                                     # ∝ 2 sinh(u)
        den = (ep + en).sum(dim=-1, keepdim=True)         # ∝ 2 Σ cosh(u)
        A = F.normalize(self.codebook, dim=-1)
        return (num @ A) / den

    @torch.no_grad()
    def export_codebook(self) -> torch.Tensor:
        return F.normalize(self.codebook.detach(), dim=-1).clone()


class SquaredReLU(nn.Module):
    def forward(self, x):
        return F.relu(x) ** 2


class RelayPatch2D(nn.Module):
    """Per-block residual relay for diffusion token streams.
    forward(x[B,T,d]) = x + sigmoid(gate) * consume(m_hat(proj(x) slots)).
    Toggle law: enabled=False returns x UNTOUCHED (code-path skip, bit-exact)."""

    def __init__(self, d: int, n_slots: int = 16, K: int = 64, tau: float = 0.1,
                 hidden: int = 178, init: str = "random"):
        super().__init__()
        self.d, self.n_slots, self.hidden = d, n_slots, hidden
        self.proj = nn.Linear(d, n_slots * 4, bias=False)
        nn.init.orthogonal_(self.proj.weight)
        self.addr = AlephAddress(K, 4, tau, init)
        self.consume = nn.Sequential(
            nn.Linear(n_slots * 4, hidden), SquaredReLU(),
            nn.LayerNorm(hidden), nn.Linear(hidden, d))
        nn.init.zeros_(self.consume[-1].weight)   # zero-init out: weight
        nn.init.zeros_(self.consume[-1].bias)     # AND bias (exp006 law)
        self.gate = nn.Parameter(torch.tensor(-3.0))
        self.enabled = True                       # plain attr, not a buffer
        self.telemetry = False                    # read-only ratio stash switch
        self.last_delta_ratio = None

    def assert_zero_init(self):
        w, b = self.consume[-1].weight, self.consume[-1].bias
        assert w.abs().max().item() == 0.0, "out weight not zero-init"
        assert b.abs().max().item() == 0.0, "out BIAS not zero-init (leak law)"

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not self.enabled:
            return x
        lead = x.shape[:-1]
        slots = self.proj(x).view(*lead, self.n_slots, 4)
        feats = self.addr.m_hat(slots).reshape(*lead, self.n_slots * 4)
        delta = torch.sigmoid(self.gate) * self.consume(feats)
        if self.telemetry:
            with torch.no_grad():
                self.last_delta_ratio = (
                    delta.norm() / x.norm().clamp_min(1e-12)).item()
        return x + delta

    @classmethod
    def from_state_dict(cls, sd: dict) -> "RelayPatch2D":
        """Rebuild dims FROM the state dict (runner_hook discipline: never trust
        constructor defaults), then strict-load."""
        d = sd["proj.weight"].shape[1]
        n_slots = sd["proj.weight"].shape[0] // 4
        hidden = sd["consume.0.weight"].shape[0]
        K = sd["addr.codebook"].shape[0]
        m = cls(d, n_slots=n_slots, K=K, hidden=hidden)
        m.load_state_dict(sd, strict=True)
        return m


class BlockWithRelay(nn.Module):
    """Wrap pattern (pod/v35_substrate.py:182 lineage): adapter reads the
    hidden-states output of the wrapped block; tuple outputs pass through."""

    def __init__(self, block: nn.Module, relay: nn.Module):
        super().__init__()
        self.block = block
        self.relay = relay

    def forward(self, *args, **kwargs):
        out = self.block(*args, **kwargs)
        if isinstance(out, tuple):
            return (self.relay(out[0]),) + out[1:]
        return self.relay(out)


class AlephCondAdapter(nn.Module):
    """Frozen [N_ADDR, addr_dim] address -> N_ADDR conditioning positions.
    Everything zero-init => tokens are EXACT zeros at init and when disabled
    (the position-presence offset is handled by the bed: SD15 masked-append,
    Anima write-into-padding)."""

    def __init__(self, cond_dim: int, addr_dim: int = 128, n_addr: int = 32):
        super().__init__()
        self.cond_dim, self.addr_dim, self.n_addr = cond_dim, addr_dim, n_addr
        self.addr_proj = nn.Linear(addr_dim, cond_dim)
        nn.init.zeros_(self.addr_proj.weight)
        nn.init.zeros_(self.addr_proj.bias)
        self.pos_table = nn.Parameter(torch.zeros(n_addr, cond_dim))
        self.gate = nn.Parameter(torch.tensor(-3.0))
        self.enabled = True

    def assert_zero_init(self):
        assert self.addr_proj.weight.abs().max().item() == 0.0
        assert self.addr_proj.bias.abs().max().item() == 0.0
        assert self.pos_table.abs().max().item() == 0.0

    def tokens(self, addr: torch.Tensor) -> torch.Tensor:
        """addr [B, n_addr, addr_dim] -> [B, n_addr, cond_dim]."""
        assert addr.shape[-2:] == (self.n_addr, self.addr_dim), addr.shape
        if not self.enabled:
            return addr.new_zeros(*addr.shape[:-1], self.cond_dim)
        return torch.sigmoid(self.gate) * (self.addr_proj(addr) + self.pos_table)

    @classmethod
    def from_state_dict(cls, sd: dict) -> "AlephCondAdapter":
        cond_dim, addr_dim = sd["addr_proj.weight"].shape
        n_addr = sd["pos_table"].shape[0]
        m = cls(cond_dim, addr_dim=addr_dim, n_addr=n_addr)
        m.load_state_dict(sd, strict=True)
        return m


# ── parity gates ─────────────────────────────────────────────────────────────

def toggle_parity(frozen_fn, adapted_fn, probes, exact: bool = True) -> float:
    """P-TOGGLE / P-INIT: max |frozen(x) - adapted(x)| over probes.
    exact=True asserts BIT-EXACT (0.0)."""
    worst = 0.0
    with torch.no_grad():
        for p in probes:
            a, b = frozen_fn(p), adapted_fn(p)
            d = (a - b).abs().max().item()
            worst = max(worst, d)
    if exact:
        assert worst == 0.0, f"parity broken: max|delta| = {worst}"
    return worst


class FireCounter:
    """P-FIRE: forward-hook counters — every adapter fires, none silently dead."""

    def __init__(self, adapters):
        self.adapters = list(adapters)
        self.counts = [0] * len(self.adapters)
        self._handles = []

    def __enter__(self):
        for i, a in enumerate(self.adapters):
            def mk(i):
                def hook(mod, inp, out):
                    self.counts[i] += 1
                return hook
            self._handles.append(a.register_forward_hook(mk(i)))
        return self

    def __exit__(self, *exc):
        for h in self._handles:
            h.remove()

    def assert_all_fired(self, at_least: int = 1):
        dead = [i for i, c in enumerate(self.counts) if c < at_least]
        assert not dead, f"adapters never fired at sites {dead}"


def save_relay_stack(adapters: nn.ModuleList, path: str):
    torch.save({"relays": [a.state_dict() for a in adapters]}, path)


def load_relay_stack(path: str) -> nn.ModuleList:
    ck = torch.load(path, map_location="cpu", weights_only=True)
    return nn.ModuleList(RelayPatch2D.from_state_dict(sd) for sd in ck["relays"])


# ── read-only gauges (argmax/pooling legal ONLY here) ────────────────────────

SIGMA_BANDS = ((1.00, 0.75), (0.75, 0.50), (0.50, 0.25), (0.25, 0.00))


def sigma_band(sigma: float) -> int:
    for i, (hi, lo) in enumerate(SIGMA_BANDS):
        if lo < sigma <= hi or (i == len(SIGMA_BANDS) - 1 and sigma <= hi):
            return i
    return 0


@torch.no_grad()
def sign_codes(relay: RelayPatch2D, x: torch.Tensor) -> torch.Tensor:
    """Per-position per-slot code: winner half-axis id*2 + orientation bit,
    read from the addressing cosines (read-only hook; gauge, not gradient)."""
    lead = x.shape[:-1]
    slots = relay.proj(x).view(*lead, relay.n_slots, relay.addr.D)
    A = F.normalize(relay.addr.codebook, dim=-1)
    cos = F.normalize(slots, dim=-1) @ A.transpose(-1, -2)
    idx = cos.abs().argmax(dim=-1)
    orient = torch.gather(cos, -1, idx.unsqueeze(-1)).squeeze(-1) > 0
    return idx * 2 + orient.long()                     # (..., n_slots)


@torch.no_grad()
def separation_no_diag(codes_by_register: dict[str, torch.Tensor]) -> float:
    """Register separation = mean inter-register Hamming − mean intra-register
    Hamming, intra DIAGONAL EXCLUDED (the r2 gauge restart; qwen-line numbers
    are NOT comparable — canon/register_probe_gauge.md). Inter keeps same-input
    (i,i) cross-register pairs: same input under two registers is a REAL pair
    (pure register effect), so two near-identical registers read sep ~ -intra/n,
    not 0. codes: (n_inputs, ..., n_slots) int tensors, inputs ALIGNED across
    registers (assert same n)."""
    regs = list(codes_by_register)
    ns = {codes_by_register[r].shape[0] for r in regs}
    assert len(ns) == 1, f"unaligned register inputs: {ns}"
    n = ns.pop()
    assert n >= max(8, n // 2), "too few aligned inputs"

    def ham(a, b):  # fraction of differing code entries
        return (a != b).float().mean().item()

    intra, inter = [], []
    for r in regs:
        c = codes_by_register[r]
        for i in range(n):
            for j in range(i + 1, n):                  # i<j: diagonal excluded
                intra.append(ham(c[i], c[j]))
    for a_i in range(len(regs)):
        for b_i in range(a_i + 1, len(regs)):
            ca, cb = codes_by_register[regs[a_i]], codes_by_register[regs[b_i]]
            for i in range(n):
                for j in range(n):
                    inter.append(ham(ca[i], cb[j]))
    return (sum(inter) / max(len(inter), 1)) - (sum(intra) / max(len(intra), 1))


def derangement(n: int, seed: int = 0) -> torch.Tensor:
    """Permutation with NO fixed points (the shuffled-control law)."""
    g = torch.Generator().manual_seed(seed)
    while True:
        p = torch.randperm(n, generator=g)
        if not (p == torch.arange(n)).any():
            return p


@torch.no_grad()
def gate_stats(adapters) -> dict:
    gates = [torch.sigmoid(a.gate).item() for a in adapters if hasattr(a, "gate")]
    return {"gate_mean": round(sum(gates) / max(len(gates), 1), 5),
            "gate_min": round(min(gates), 5) if gates else None,
            "gate_max": round(max(gates), 5) if gates else None}


# ── CPU cert suite (shapes/parity ONLY — never CPU-train) ───────────────────

def _smoke():
    torch.manual_seed(0)

    # 1) super-fibonacci: shape + unit norm
    q = super_fibonacci_s3(64)
    assert q.shape == (64, 4)
    assert (q.norm(dim=-1) - 1).abs().max() < 1e-6

    # 2) closed-form m_hat == explicit 2K softmax read (<=1e-6 fp32)
    addr = AlephAddress(K=64, D=4)
    x = torch.randn(5, 7, 16, 4)
    u = addr._u(x)
    p = torch.softmax(torch.cat([u, -u], dim=-1), dim=-1)
    A = F.normalize(addr.codebook, dim=-1)
    explicit = p[..., :64] @ A - p[..., 64:] @ A
    assert (addr.m_hat(x) - explicit).abs().max() < 1e-6, "closed form != 2K softmax"

    # 3) RelayPatch2D: zero-init, init parity, toggle parity, shape
    for d in (320, 640, 1280, 2048):
        r = RelayPatch2D(d)
        r.assert_zero_init()
        xb = torch.randn(2, 9, d)
        assert torch.equal(r(xb), xb), "P-INIT broken (zero-init must be exact)"
        r.enabled = False
        assert r(xb) is xb, "P-TOGGLE broken (must be code-path skip)"
        r.enabled = True
        nn.init.normal_(r.consume[-1].weight, std=0.02)   # pretend trained
        nn.init.normal_(r.consume[-1].bias, std=0.02)
        assert r(xb).shape == xb.shape and not torch.equal(r(xb), xb)
    n_params = sum(p.numel() for p in RelayPatch2D(1280).parameters())
    print(f"  relay params @d=1280: {n_params:,}")

    # 4) toy-stub toggle parity through BlockWithRelay
    class ToyBlock(nn.Module):
        def __init__(self, d):
            super().__init__()
            self.lin = nn.Linear(d, d)

        def forward(self, x):
            return x + torch.tanh(self.lin(x))

    d = 64
    torch.manual_seed(1)
    frozen = nn.Sequential(ToyBlock(d), ToyBlock(d))
    relays = nn.ModuleList([RelayPatch2D(d) for _ in range(2)])
    for rl in relays:                                  # pretend trained
        nn.init.normal_(rl.consume[-1].weight, std=0.02)
    wrapped = nn.Sequential(*[BlockWithRelay(b, r)
                              for b, r in zip(frozen, relays)])
    probes = [torch.randn(2, 5, d) for _ in range(4)]
    for rl in relays:
        rl.enabled = False
    toggle_parity(lambda p: frozen(p), lambda p: wrapped(p), probes, exact=True)
    for rl in relays:
        rl.enabled = True
    worst = toggle_parity(lambda p: frozen(p), lambda p: wrapped(p), probes,
                          exact=False)
    assert worst > 0, "trained adapters produced no delta (dead sites?)"

    # 5) P-FIRE
    with FireCounter(relays) as fc:
        wrapped(probes[0])
    fc.assert_all_fired()

    # 6) P-SAVE: dims-from-state-dict round trip + re-parity
    import tempfile, os
    path = os.path.join(tempfile.gettempdir(), "r2_relay_smoke.pt")
    save_relay_stack(relays, path)
    relays2 = load_relay_stack(path)
    for a, b in zip(relays, relays2):
        for (ka, va), (kb, vb) in zip(a.state_dict().items(),
                                      b.state_dict().items()):
            assert ka == kb and torch.equal(va, vb)

    # 7) AlephCondAdapter: exact zeros at init + disabled + round trip
    for cd in (768, 1024, 2048):
        c = AlephCondAdapter(cd)
        c.assert_zero_init()
        a32 = torch.randn(3, 32, 128)
        assert c.tokens(a32).abs().max().item() == 0.0, "cond tokens not exact 0"
        c.enabled = False
        assert c.tokens(a32).abs().max().item() == 0.0
        c2 = AlephCondAdapter.from_state_dict(c.state_dict())
        assert c2.cond_dim == cd

    # 8) sign codes + separation gauge (intra diagonal excluded)
    r = RelayPatch2D(64)
    xa = torch.randn(10, 6, 64)
    codes_a = sign_codes(r, xa)
    assert codes_a.shape == (10, 6, 16)
    assert torch.equal(codes_a, sign_codes(r, xa)), "codes not deterministic"
    # mechanics on constructed codes: intra of constant register = 0 (diagonal
    # exclusion working), all-differ registers => inter 1 => sep exactly 1.0
    cz = torch.zeros(10, 6, 16, dtype=torch.long)
    co = torch.ones(10, 6, 16, dtype=torch.long)
    assert separation_no_diag({"z": cz, "o": co}) == 1.0
    assert separation_no_diag({"z": cz, "z2": cz.clone()}) == 0.0
    # near-identical REAL registers read ~ -intra/n (documented, not a bug)
    same = separation_no_diag({"a": codes_a, "b": codes_a.clone()})
    assert same <= 0.0
    print(f"  sep(identical real)={same:.4f} (expected ~ -intra/n)")

    # 9) derangement: no fixed points
    for n in (5, 12, 24):
        p = derangement(n, seed=3)
        assert not (p == torch.arange(n)).any()

    # 10) sigma bands cover [0,1]
    assert [sigma_band(s) for s in (1.0, 0.8, 0.6, 0.3, 0.0)] == [0, 0, 1, 2, 3]

    print("aleph_diffusion_core smoke PASSED (closed-form parity, zero-init "
          "W+b, bit-exact toggles, fire counters, save/load dims-from-sd, "
          "cond-adapter exact zeros, no-diag gauge, derangement)")


if __name__ == "__main__":
    _smoke()