src/tests/test_bind_invariants.py

"""E3-BIND invariants. Preregistered as gating: Stage A may not be reported
unless every one of these passes."""
import sys
from pathlib import Path

import numpy as np
import pytest
import torch

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from pns.data.binding import N_BIND_SLOTS, binding_plan
from pns.model.bind import BindConfig, PNSBind
from pns.world.schema import ET, Fam
from pns.world.sim import generate_lifetime


def test_unrelated_binding_bitwise_unchanged():
    """The defining intervention: an event writes its own slot and NOTHING
    else, bitwise, across thousands of randomised events."""
    torch.manual_seed(0)
    m = PNSBind(BindConfig(mode="bind")).eval()
    B, S, d = 16, N_BIND_SLOTS, m.cfg.d
    rng = np.random.default_rng(0)
    total_checked = 0
    with torch.no_grad():
        for it in range(40):                       # 40 x 16 = 640 events...
            state = torch.randn(B, S, d)
            tok = torch.randint(1, 8000, (B, 48))
            tok[:, 40:] = 0
            w = torch.from_numpy(rng.integers(-1, S, B)).long()
            r = torch.from_numpy(rng.integers(-1, S, B)).long()
            se = torch.from_numpy(rng.integers(0, 200, (B, S))).short()
            sa = torch.from_numpy(rng.integers(0, 14, (B, S))).short()
            new, _ = m.step(state, tok, torch.randint(0, len(ET), (B,)),
                            torch.randint(1, 60, (B,)), w, r, se, sa)
            for b in range(B):
                wi = int(w[b])
                mask = torch.ones(S, dtype=torch.bool)
                if wi >= 0:
                    mask[wi] = False
                # every non-addressed slot must be bitwise identical
                assert torch.equal(new[b][mask], state[b][mask]), \
                    f"iter {it} row {b}: unrelated slots changed"
                if wi >= 0:
                    assert not torch.equal(new[b, wi], state[b, wi]), \
                        "the addressed slot did not change"
                total_checked += int(mask.sum())
    assert total_checked > 100_000, total_checked


def test_address_contains_no_semantic_value():
    """The exact side may carry identity and version only - never the value."""
    for seed in range(40):
        lt = generate_lifetime(seed, 256)
        bp = binding_plan(lt)
        # the only exact per-slot fields are entity id and attribute type
        assert set(bp) >= {"write_slot", "read_slot", "slot_ent", "slot_attr",
                           "slot_ver"}
        vals = set()
        for ev in lt["events"]:
            v = ev.payload.get("value")
            if isinstance(v, str):
                vals.add(v)
        # no packed exact field may equal any semantic value (they are ints)
        for k in ("slot_ent", "slot_attr", "slot_ver", "write_slot", "read_slot"):
            assert bp[k].dtype.kind in "iu", k


def test_slot_allocation_independent_of_label():
    """Slot index must not encode the answer: permuting values leaves the
    allocation identical, because allocation is seeded by lifetime id only."""
    for seed in range(30):
        lt = generate_lifetime(seed, 256)
        a = binding_plan(lt)
        # mutate every semantic value; allocation must not move
        lt2 = generate_lifetime(seed, 256)
        for ev in lt2["events"]:
            if isinstance(ev.payload.get("value"), str):
                ev.payload["value"] = "ZZZ-MUTATED"
        b = binding_plan(lt2)
        assert np.array_equal(a["write_slot"], b["write_slot"])
        assert np.array_equal(a["read_slot"], b["read_slot"])
        assert np.array_equal(a["slot_ent"], b["slot_ent"])


def test_no_unregistered_slot_reuse():
    """No slot is ever allocated twice, and capacity is never exceeded."""
    worst = 0
    for seed in range(300):
        lt = generate_lifetime(seed, 256)
        bp = binding_plan(lt)
        used = [s for s in bp["write_slot"].tolist() if s >= 0]
        assert len(used) == len(set(used)), f"slot reused in lifetime {seed}"
        assert bp["overflow"] == 0, f"capacity overflow in lifetime {seed}"
        worst = max(worst, bp["n_used"])
    assert worst < N_BIND_SLOTS, worst


def test_revert_changes_authority_pointer_not_payload():
    """A revert must allocate no slot and write no payload - it only moves the
    active pointer to an earlier version's slot."""
    found = 0
    for seed in range(200):
        lt = generate_lifetime(seed, 256)
        bp = binding_plan(lt)
        for i, ev in enumerate(lt["events"]):
            if ev.etype == ET.CORRECTION and ev.payload.get("mode") == "revert":
                assert bp["write_slot"][i] == -1, \
                    "a revert allocated a slot / wrote a payload"
                found += 1
    assert found > 20, f"too few reverts exercised ({found})"


def test_read_slot_resolves_for_semantic_questions():
    """Sanity: the addressed read actually resolves for the queried family."""
    tot = hit = 0
    for seed in range(120):
        lt = generate_lifetime(seed, 256)
        bp = binding_plan(lt)
        for i, ev in enumerate(lt["events"]):
            if ev.etype == ET.QUESTION and ev.family in (int(Fam.SEM_LATEST),
                                                         int(Fam.SEM_2HOP)):
                tot += 1
                hit += int(bp["read_slot"][i] >= 0)
    assert tot > 200 and hit / tot > 0.9, (hit, tot)


def test_read_path_identical_between_arms():
    """CRITICAL for E3-A3: Bind and Unbound must differ ONLY in the write
    transition. Both receive the same exact address at query time. If Bind had
    a privileged retrieval path, a positive result would conflate 'protected
    writing' with 'easier retrieval'.

    Executable form: with writes frozen, identical weights and identical
    inputs, the two modes must produce BITWISE identical outputs."""
    torch.manual_seed(7)
    a = PNSBind(BindConfig(mode="bind")).eval()
    torch.manual_seed(7)
    b = PNSBind(BindConfig(mode="unbound")).eval()
    # same seed must give identical initial parameters (pairing precondition)
    for (na, pa), (nb, pb) in zip(a.named_parameters(), b.named_parameters()):
        assert na == nb and torch.equal(pa, pb), na

    B, S, d = 8, a.cfg.n_slots, a.cfg.d
    rng = np.random.default_rng(3)
    with torch.no_grad():
        for _ in range(25):
            state = torch.randn(B, S, d)
            tok = torch.randint(1, 8000, (B, 48))
            tok[:, 40:] = 0
            w = torch.from_numpy(rng.integers(-1, S, B)).long()
            r = torch.from_numpy(rng.integers(-1, S, B)).long()
            se = torch.from_numpy(rng.integers(0, 200, (B, S))).short()
            sa = torch.from_numpy(rng.integers(0, 14, (B, S))).short()
            et = torch.randint(0, len(ET), (B,))
            dt = torch.randint(1, 60, (B,))
            _, oa = a.step(state, tok, et, dt, w, r, se, sa, freeze_writes=True)
            _, ob = b.step(state, tok, et, dt, w, r, se, sa, freeze_writes=True)
            assert torch.equal(oa["enum"], ob["enum"]), \
                "read paths differ between arms - E3-A3 comparison invalid"


def test_arms_differ_only_in_write():
    """The converse: with writes ENABLED the two arms must diverge, otherwise
    the intervention is not actually being applied."""
    torch.manual_seed(11)
    a = PNSBind(BindConfig(mode="bind")).eval()
    torch.manual_seed(11)
    b = PNSBind(BindConfig(mode="unbound")).eval()
    B, S, d = 4, a.cfg.n_slots, a.cfg.d
    state = torch.randn(B, S, d)
    tok = torch.randint(1, 8000, (B, 48))
    w = torch.zeros(B, dtype=torch.long)
    r = torch.zeros(B, dtype=torch.long)
    se = torch.zeros(B, S, dtype=torch.short)
    sa = torch.zeros(B, S, dtype=torch.short)
    with torch.no_grad():
        sa_, _ = a.step(state, tok, torch.zeros(B, dtype=torch.long),
                        torch.ones(B, dtype=torch.long), w, r, se, sa)
        sb_, _ = b.step(state, tok, torch.zeros(B, dtype=torch.long),
                        torch.ones(B, dtype=torch.long), w, r, se, sa)
    assert not torch.equal(sa_, sb_), "arms did not diverge under writes"
    # and specifically: bind left slot 1 alone, unbound did not
    assert torch.equal(sa_[:, 1], state[:, 1])
    assert not torch.equal(sb_[:, 1], state[:, 1])