# SPDX-License-Identifier: BSD-3-Clause-Clear
# Copyright (c) 2026 Ronald Lee Hampton. All rights reserved.
# No express or implied licenses to any party's patent rights are granted by this license.
#
# read_shuffle_cost.py: re-measures the shuffle cost of every shipped brain and re-counts the switchboards.
#
# Usage (from the package folder):
#   python -I reader/read_shuffle_cost.py --brains brains --data data --expect results/per_brain_record.csv
# Needs: Python 3, numpy, torch, safetensors. CPU only, one thread, about a minute for the 80 brains. No network.
# Side effects: none; it prints only (redirect the output to keep it).
#
# THE MEASURE. Each brain is run on the 256 held-out windows (66 tokens each: two context words, then 64 targets), 16,384 positions in
# one batch. The reasoning core's quantized output (qk, one 32-number vector per position) is then shuffled across positions with
# numpy default_rng(0).permutation(16384), everything else held fixed, and the network is run again. The SHUFFLE COST is the mean, over
# the positions whose target is not one of the 40 excluded ids, of (cross-entropy shuffled minus cross-entropy as trained), in nats.
# THE LINE (fixed before the run): a core is a SWITCHBOARD if its shuffle cost is below 0.025 nats, else a CONTEXT WORKER.
# A RESCUE is a seed whose plain (ST-0) core is a switchboard and whose sleep-form (S-SLEEP) core is a context worker; a HARM is the
# reverse. THE BAR (fixed before the run): S-SLEEP ends with fewer switchboards than ST-0, AND at least 5 seeds are discordant, AND the
# exact one-sided sign test on the discordant seeds gives p at or below 0.05.
#
# This file holds only the forward pass the measure needs (no training code). Tensor names keep the run's own prefixes:
#   tiny.*  the routed base network: two word embeddings summed, a router to one of 48 experts (top 1, its score nudged by the pull
#           0.4 x (mean size minus the expert's size), sizes in "y"), each expert a two-layer block, a layer norm and a shared head;
#   rc.*    the reasoning core: written from the integer part's hidden state, one gated pass, quantized to whole numbers in -127..127;
#   cc.*    the small integer part with ternary weights: reads the two context words' ids as 12 bits each, and writes a gain on the
#           base network's expert output from its own hidden state and the core's quantized output (zero gain at the start).
#   cc.Wv and cc.bv, a two-output head the measure does not use, were removed from the shipped files (results/head_removed_record.csv).
# Every operation below is written in the same order as the run's own code so that the numbers can match it bit for bit on the same
# software and hardware; on other hardware or torch versions the last digits may differ (the --expect check reports by how much).
import sys, os, csv, glob, math, hashlib, argparse
import numpy as np
import torch
import torch.nn.functional as F

V, D, T, CTX, NBITS = 4096, 32, 64, 2, 12
S_PULL, RHO = 0.4, 0.5
LINE = 0.025
FORMS = ("ST-0", "S-SLEEP")
P_BAR, MIN_DISC = 0.05, 5
BITW = 2 ** torch.arange(NBITS)


def sha256_file(p):
    h = hashlib.sha256()
    with open(p, "rb") as f:
        for blk in iter(lambda: f.read(1 << 20), b""):
            h.update(blk)
    return h.hexdigest()


def sign_p(r, n):
    """exact one-sided p of at least r rescues among n discordant seeds, each a fair coin."""
    return float(sum(math.comb(n, k) for k in range(r, n + 1)) / 2 ** n) if n else 1.0


def tern(w):
    return (w > 0.5).float() - (w < -0.5).float()


def ste_tern(w):                           # the training-time expression, kept so the floating-point result is the same
    return w + (tern(w) - w)


def ste_round(x):
    return x + (torch.round(x) - x)


def ste_floor(x):
    return x + (torch.floor(x) - x)


def rms(x, w):
    return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * w


def msg(xs):
    """the two context words' token ids as 12 bits each, +1/-1, newest word first."""
    parts = [((x[:, None] // BITW) % 2).float() * 2 - 1 for x in xs[:2]]
    return torch.cat(parts, dim=-1)


def hidden(P, xs):
    """the integer part's hidden state: ternary weights, whole-number biases, clamps; every value a whole number."""
    m = msg(xs)
    a1 = m @ ste_tern(P["cc.W1"]).t() + ste_round(P["cc.b1"] * 1.0)
    h1 = torch.clamp(a1, 0.0, 31.0)
    a2 = h1 @ ste_tern(P["cc.W2"]).t() + ste_round(P["cc.b2"] * 16.0)
    return torch.clamp(ste_floor(a2 / 4.0), 0.0, 127.0)


def pull_bias(P):
    y = P["y"].numpy()
    return torch.tensor(S_PULL * (y.mean() - y), dtype=torch.float32)


def forward(P, xs, bias, h2, perm=None):
    """logits at every position; with perm, the reasoning core's quantized output is shuffled across positions first."""
    # the routed base network
    x = F.embedding(xs[0], P["tiny.emb.0.weight"])
    x = x + F.embedding(xs[1], P["tiny.emb.1.weight"])
    r = F.linear(x, P["tiny.router.weight"], P["tiny.router.bias"])
    score = r + bias
    p = F.softmax(r, dim=-1)
    sel = torch.argmax(score, dim=-1)
    g = p.gather(1, sel[:, None]).squeeze(1)
    h = F.gelu(torch.bmm(x[:, None, :], P["tiny.w1"][sel])).squeeze(1)
    y = torch.bmm(h[:, None, :], P["tiny.w2"][sel]).squeeze(1)
    contrib = g[:, None] * y
    # the reasoning core, written from the integer part's hidden state
    K0 = (1 - RHO) * torch.tanh(rms(h2 / 127.0, P["rc.nwc.weight"]) @ P["rc.wkc"])
    z = rms(K0, P["rc.ns.weight"])
    u = F.silu(z @ P["rc.w1"]) * (z @ P["rc.w3"])
    K2 = RHO * K0 + (1 - RHO) * torch.tanh(u @ P["rc.wh"])
    qk = torch.clamp(ste_round(K2 * 127.0), -127.0, 127.0)
    if perm is not None:
        qk = qk[perm]                      # THE SHUFFLE: the core's output moved across positions, nothing else changed
    # the integer part's write: a gain on the expert output
    c = h2 @ ste_tern(P["cc.Wcg"]).t()
    c = c + qk @ ste_tern(P["cc.Wck"]).t()
    gain = 2.0 * torch.sigmoid(c / 256.0) - 1.0
    o = x + contrib * (1.0 + gain)
    o = F.layer_norm(o, (D,), P["tiny.norm.weight"], P["tiny.norm.bias"], 1e-5)
    return F.linear(o, P["tiny.head.weight"], P["tiny.head.bias"])


def heldout(data_dir):
    win = np.load(os.path.join(data_dir, "heldout_windows_uint16.npy"), allow_pickle=False).astype(np.int64)
    if win.shape != (256, T + CTX):
        raise SystemExit("REFUSED: held-out windows have shape %s" % (win.shape,))
    xs = [torch.from_numpy(np.ascontiguousarray(win[:, CTX - 1 - k:T + CTX - 1 - k])).reshape(-1) for k in range(CTX)]
    tgt = torch.from_numpy(np.ascontiguousarray(win[:, CTX:])).reshape(-1)
    ids = [int(l) for l in open(os.path.join(data_dir, "excluded_target_ids.txt"), encoding="utf-8") if l.strip() and not l.startswith("#")]
    excl = torch.zeros(V, dtype=torch.bool)
    excl[torch.tensor(ids)] = True
    allowed = (~excl[tgt]).numpy()
    return win, xs, tgt, allowed, ids


def main(a):
    from safetensors.torch import load_file
    torch.set_num_threads(1)
    win, xs, tgt, allowed, ids = heldout(a.data)
    perm = torch.from_numpy(np.random.default_rng(0).permutation(len(tgt)))
    print("held-out windows %d, positions %d, scored positions %d, excluded ids %d" % (win.shape[0], len(tgt), int(allowed.sum()), len(ids)))
    print("torch %s, numpy %s, threads %d" % (torch.__version__, np.__version__, torch.get_num_threads()))
    paths = sorted(glob.glob(os.path.join(a.brains, "*.safetensors")))
    per = {}
    for path in paths:
        name = os.path.basename(path)[:-len(".safetensors")]
        form, s = name.split("_")[0], int(name.split("_")[1][1:])
        P = load_file(path)
        with torch.no_grad():
            bias = pull_bias(P)
            h2 = hidden(P, xs)
            lg_on = forward(P, xs, bias, h2)
            lg_sh = forward(P, xs, bias, h2, perm)
        ce_on = F.cross_entropy(lg_on, tgt, reduction="none").double().numpy()
        ce_sh = F.cross_entropy(lg_sh, tgt, reduction="none").double().numpy()
        cost = float((ce_sh - ce_on)[allowed].mean())
        per[(form, s)] = {"name": name, "cost": cost, "label": "SWITCHBOARD" if cost < LINE else "CONTEXT WORKER",
                          "ce_on": float(ce_on[allowed].mean()),
                          "logits_sha256": hashlib.sha256(lg_on.contiguous().numpy().tobytes()).hexdigest(),
                          "file_sha256": sha256_file(path)}
        print("%-8s s%d  shuffle cost %.6f  %-14s  (held-out loss as trained %.6f)" % (form, s, cost, per[(form, s)]["label"], per[(form, s)]["ce_on"]))
    seeds = sorted({s for (_f, s) in per})
    if not seeds or not all((f, s) in per for f in FORMS for s in seeds):
        print("REFUSED: no brains, or a seed is missing one of its two brains")
        return 2
    sb = {f: [s for s in seeds if per[(f, s)]["label"] == "SWITCHBOARD"] for f in FORMS}
    res = [s for s in seeds if per[("ST-0", s)]["label"] == "SWITCHBOARD" and per[("S-SLEEP", s)]["label"] == "CONTEXT WORKER"]
    harm = [s for s in seeds if per[("ST-0", s)]["label"] == "CONTEXT WORKER" and per[("S-SLEEP", s)]["label"] == "SWITCHBOARD"]
    n = len(res) + len(harm)
    p = sign_p(len(res), n)
    passed = len(sb["S-SLEEP"]) < len(sb["ST-0"]) and n >= MIN_DISC and p <= P_BAR
    print("")
    print("seeds %d (%d to %d); THE LINE: switchboard below %.3f nats of shuffle cost" % (len(seeds), seeds[0], seeds[-1], LINE))
    for f in FORMS:
        print("  %-8s switchboards %d of %d: %s" % (f, len(sb[f]), len(seeds), sb[f]))
    print("  rescues %d %s; harms %d %s; discordant %d; exact one-sided sign test p %.5f" % (len(res), res, len(harm), harm, n, p))
    print("THE BAR (S-SLEEP fewer, at least %d discordant, p <= %.2f): %s" % (MIN_DISC, P_BAR, "CONFIRMED" if passed else "NOT CONFIRMED"))
    near = sorted((abs(v["cost"] - LINE), "%s s%d" % k, round(v["cost"], 6)) for k, v in per.items())[:3]
    print("  nearest the line: %s" % [(k, c) for _d, k, c in near])
    rc = 0
    if a.expect:
        rows = {(r["form"], int(r["seed"])): r for r in csv.DictReader(open(a.expect, encoding="utf-8", newline=""))}
        keys = [k for k in per if k in rows]
        lab_eq = sum(1 for k in keys if rows[k]["label"] == per[k]["label"])
        file_eq = sum(1 for k in keys if rows[k]["safetensors_sha256"] == per[k]["file_sha256"])
        cost_eq = sum(1 for k in keys if float(rows[k]["shuffle_cost"]) == per[k]["cost"])
        ce_eq = sum(1 for k in keys if float(rows[k]["ce_on_allowed_mean"]) == per[k]["ce_on"])
        lg_eq = sum(1 for k in keys if rows[k]["logits_sha256_on"] == per[k]["logits_sha256"])
        dmax = max(abs(float(rows[k]["shuffle_cost"]) - per[k]["cost"]) for k in keys) if keys else float("nan")
        cmax = max(abs(float(rows[k]["ce_on_allowed_mean"]) - per[k]["ce_on"]) for k in keys) if keys else float("nan")
        print("")
        print("AGAINST %s (%d rows):" % (os.path.basename(a.expect), len(rows)))
        print("  brain files equal to the record's sha256: %d of %d" % (file_eq, len(per)))
        print("  labels equal: %d of %d" % (lab_eq, len(per)))
        print("  shuffle costs equal bit for bit: %d of %d; largest absolute difference %.3g nats" % (cost_eq, len(per), dmax))
        print("  held-out loss as trained equal bit for bit: %d of %d; largest absolute difference %.3g nats" % (ce_eq, len(per), cmax))
        print("  logits as trained equal bit for bit (sha256): %d of %d" % (lg_eq, len(per)))
        if lab_eq != len(per) or file_eq != len(per) or len(rows) != len(per):
            rc = 1
    return rc


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--brains", default="brains")
    ap.add_argument("--data", default="data")
    ap.add_argument("--expect", default="")
    sys.exit(main(ap.parse_args()))