"""The reviewer's decisive diagnostic, run against the model the trader actually loads.

Question 1 (batch-composition leakage): does one sample's output change when the OTHER
samples in its batch change? If yes, batch statistics are leaking and a sample's
representation depends on samples that may be from its future.
Question 2 (train/serve asymmetry): does B=1 differ from the same sample inside B=64?

Weights are random on purpose: batch coupling is a property of the ARCHITECTURE, not of
the weight values, and this avoids loading 5.9GB next to the live trader's VRAM.
"""
import sys, numpy as np, torch
sys.path.insert(0, "/home/kairosliquid")
import dreadnought_byte_dag as D

torch.manual_seed(0); np.random.seed(0)
SEQ = 128
model = D.DreadnoughtByteDAG(seq_len=SEQ, dim=64, ffn_dim=128, vocab=256,
                             state_dim=64, conv_channels=32, use_whiteboard=True)
model.eval()

rng = np.random.default_rng(0)
probe = rng.integers(0, 256, size=(1, SEQ), dtype=np.int64)          # the sample under test
early = rng.integers(0, 100, size=(31, SEQ), dtype=np.int64)          # "early period" companions
late  = rng.integers(150, 256, size=(31, SEQ), dtype=np.int64)        # "late period" companions

def run(batch):
    with torch.no_grad():
        out = model(torch.tensor(batch, dtype=torch.long))
    return (out[0] if isinstance(out, (tuple, list)) else out).detach().float().numpy()

a = run(np.concatenate([probe, early]))[0]     # probe first, early companions
b = run(np.concatenate([probe, late]))[0]      # probe first, late companions
solo = run(probe)[0]                            # B=1
big  = run(np.concatenate([probe, early, late]))[0]   # B=63

def report(name, x, y):
    d = float(np.max(np.abs(x - y)))
    print(f"{name:<42} max|Δ| = {d:.3e}   {'IDENTICAL' if d == 0.0 else ('within fp noise' if d < 1e-6 else '*** DIFFERS ***')}")

print(f"model: DreadnoughtByteDAG(use_whiteboard=True), eval(), random weights, output dim {a.shape}")
report("Q1  early-companions vs late-companions", a, b)
report("Q2a B=1 vs B=32 (early companions)", solo, a)
report("Q2b B=1 vs B=63 (mixed companions)", solo, big)
print()
print("norm layers present:", sorted({type(m).__name__ for m in model.modules()
                                     if 'Norm' in type(m).__name__}))

# ── SELF-TEST: prove this diagnostic can FAIL. A test that only ever says "clean"
# is a green light with no wiring. Inject a BatchNorm1d on the encoder output —
# the exact mechanism the critique describes — and re-run Q1.
if "--selftest" in sys.argv:
    print("\n--- SELF-TEST: injecting BatchNorm1d (the leak mechanism) ---")
    enc = model.encoder
    bn = torch.nn.BatchNorm1d(64)
    bn.train()  # batch statistics, as in training
    class Leaky(torch.nn.Module):
        def __init__(self, inner, bn): super().__init__(); self.inner=inner; self.bn=bn
        def forward(self, x): return self.bn(self.inner(x))
    model.encoder = Leaky(enc, bn)
    a2 = run(np.concatenate([probe, early]))[0]
    b2 = run(np.concatenate([probe, late]))[0]
    d = float(np.max(np.abs(a2 - b2)))
    print(f"Q1 with BatchNorm injected                 max|Δ| = {d:.3e}   "
          f"{'*** DIFFERS — test fires correctly ***' if d > 1e-6 else 'DID NOT FIRE (test is broken)'}")
    sys.exit(0 if d > 1e-6 else 1)