src/matrix8.py

"""neural_matrix8 — the 8-bit threshold CPU as a recurrent ternary linear-threshold
network.

The verified gate wiring of the registers-profile CPU (8-bit data, 4-bit
addresses, 16 B memory) is compiled into a fixed sequence of ternary weight
matrices W1..Wk with a Heaviside step between them, acting on a 173-bit state
vector that holds the program counter, all four registers, the flags, the
stack pointer, the halt bit, and every memory bit:

    [ PC[4] | R0[8] R1[8] R2[8] R3[8] | Z N C V | SP[4] | HALT | MEM[16][8] ]

One instruction is one pass through the stack (s' = H(Wk @ ... H(W1 @ s + b1)
... + bk)); the machine is the stack iterated until the halt bit sets, so the
whole processor is a single recurrent linear-threshold network. The transition
is total: a halted state is a fixed point (every architectural write is gated
by NOT halt), so iteration past halt is harmless.

Every weight is in {-1, 0, 1} and every bias is a small integer; the circuit
is assembled from the same verified cell constructions the family ships
(two-layer OR/NAND XOR cells, ha1/ha2/carry_or full adders, bit-cascade
comparators, 2:1 mux cells, one-threshold-gate address decode rows), then
levelized ASAP with identity pass-throughs (H(x-1)) carrying live signals
between layers.

Equality with the gate-graph processor is machine-checked, bit for bit:
  - exhaustively over all 65,536 (a, b) operand pairs for every ALU opcode
    (ADD SUB AND OR XOR SHL SHR MUL DIV CMP), batched through the matrices;
  - exhaustively over all Jcc condition x flag combinations, all JMP/CALL
    targets and SP values, and all LOAD/STORE address/data pairs;
  - by full-state lockstep against GenericThresholdCPU walking the shipped
    neural_computer8_registers weights, per instruction, on a micro-program
    suite covering every opcode;
  - by single-step agreement with a pure-integer reference on random states.

Analog realization: pre-activations are integers, 1-firing gates sit at >= 0
and 0-gates at <= -1, so placing the analog comparator threshold at -0.5 gives
every gate a symmetric guaranteed noise margin of 0.5 — any total analog error
(read noise + conductance mismatch) below 0.5 per pre-activation provably
reproduces the digital circuit bit-exactly. The analog suite measures the
margin, then sweeps Gaussian read noise and static ternary-conductance
mismatch to locate the empirical breakdown against the guarantee.

Usage:
    python matrix8.py build      # compile + save variants/neural_matrix8.safetensors
    python matrix8.py verify     # equality suite (build in memory if file absent)
    python matrix8.py analog     # margin + noise / mismatch sweeps
    python matrix8.py all        # build + verify + analog
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from typing import Dict, List, Optional, Tuple

import torch
from safetensors import safe_open
from safetensors.torch import save_file

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))  # repo root; this module lives in src/
MODEL_PATH = os.path.join(REPO, "variants", "neural_matrix8.safetensors")
GATE_MODEL_PATH = os.path.join(REPO, "variants", "neural_computer8_registers.safetensors")

ADDR_BITS = 4
MEM_BYTES = 16
STATE_BITS = 4 + 32 + 4 + 4 + 1 + 128  # pc, regs, flags, sp, halt, mem


# =============================================================================
# Circuit IR: named threshold gates over binary signals, constants fold away
# =============================================================================

class Net:
    """A DAG of threshold gates. Signals are names; '#0' / '#1' are constants
    folded into biases. Helper constructors perform algebraic collapse
    (AND(x,#1)=x, OR with one live input = that input, ...) so the compiled
    stack carries no dead logic; the collapses change nothing semantically and
    the whole machine is verified bit-for-bit afterwards."""

    def __init__(self):
        self.gates: Dict[str, Tuple[List[Tuple[str, int]], int]] = {}

    def g(self, name: str, ins: List[Tuple[str, int]], bias: int) -> str:
        folded: List[Tuple[str, int]] = []
        b = int(bias)
        for s, w in ins:
            assert w in (-1, 1), f"non-ternary weight {w} at {name}"
            if s == "#1":
                b += w
            elif s == "#0":
                pass
            else:
                folded.append((s, int(w)))
        if name in self.gates:
            raise ValueError(f"duplicate gate {name}")
        self.gates[name] = (folded, b)
        return name

    # --- verified cell constructions --------------------------------------
    def NOT(self, name, x):
        if x == "#1":
            return "#0"
        if x == "#0":
            return "#1"
        return self.g(name, [(x, -1)], 0)

    def BUF(self, name, x):
        if x in ("#0", "#1"):
            return x
        return self.g(name, [(x, 1)], -1)

    def AND(self, name, ins: List[str]):
        live = []
        for s in ins:
            if s == "#0":
                return "#0"
            if s != "#1":
                live.append(s)
        if not live:
            return "#1"
        if len(live) == 1:
            return live[0]
        return self.g(name, [(s, 1) for s in live], -len(live))

    def OR(self, name, ins: List[str]):
        live = []
        for s in ins:
            if s == "#1":
                return "#1"
            if s != "#0":
                live.append(s)
        if not live:
            return "#0"
        if len(live) == 1:
            return live[0]
        return self.g(name, [(s, 1) for s in live], -1)

    def NOR(self, name, ins: List[str]):
        live = [s for s in ins if s != "#0"]
        if any(s == "#1" for s in live):
            return "#0"
        if not live:
            return "#1"
        return self.g(name, [(s, -1) for s in live], 0)

    def XOR(self, prefix, a, b):
        """Two-layer OR/NAND XOR, the family's verified construction."""
        if a == "#0":
            return b if b in ("#0", "#1") else self.BUF(f"{prefix}.buf", b)
        if b == "#0":
            return a if a in ("#0", "#1") else self.BUF(f"{prefix}.buf", a)
        if a == "#1":
            return self.NOT(f"{prefix}.not", b)
        if b == "#1":
            return self.NOT(f"{prefix}.not", a)
        h_or = self.g(f"{prefix}.l1or", [(a, 1), (b, 1)], -1)
        h_nand = self.g(f"{prefix}.l1nand", [(a, -1), (b, -1)], 1)
        return self.g(prefix, [(h_or, 1), (h_nand, 1)], -2)

    def FA(self, prefix, a, b, cin):
        """Verified full-adder cell (ha1 XOR/AND, ha2 XOR/AND, carry OR)."""
        s1 = self.XOR(f"{prefix}.ha1s", a, b)
        c1 = self.AND(f"{prefix}.ha1c", [a, b])
        s2 = self.XOR(f"{prefix}.ha2s", s1, cin)
        c2 = self.AND(f"{prefix}.ha2c", [s1, cin])
        co = self.OR(f"{prefix}.cor", [c1, c2])
        return s2, co

    def MUX(self, prefix, sel, x1, x0):
        """2:1 mux cell: sel ? x1 : x0 (not_sel / and / and / or)."""
        if sel == "#1":
            return x1
        if sel == "#0":
            return x0
        if x1 == x0:
            return x0
        ns = self.NOT(f"{prefix}.ns", sel)
        a1 = self.AND(f"{prefix}.a1", [x1, sel])
        a0 = self.AND(f"{prefix}.a0", [x0, ns])
        return self.OR(f"{prefix}.or", [a1, a0])

    def DECODE(self, name, bits: List[str], value: int):
        """One-gate address-decode row (the memory.addr_decode construction):
        fires iff the MSB-first bits equal `value`."""
        n = len(bits)
        ins = []
        pop = 0
        for i in range(n):
            bit = (value >> (n - 1 - i)) & 1
            ins.append((bits[i], 1 if bit else -1))
            pop += bit
        return self.g(name, ins, -pop)


# =============================================================================
# State layout
# =============================================================================

def state_names() -> List[str]:
    names = [f"pc{i}" for i in range(4)]
    for r in range(4):
        names += [f"r{r}b{i}" for i in range(8)]
    names += ["fZ", "fN", "fC", "fV"]
    names += [f"sp{i}" for i in range(4)]
    names += ["halt"]
    for j in range(MEM_BYTES):
        names += [f"m{j}b{i}" for i in range(8)]
    return names


def state_to_vec(s: dict) -> torch.Tensor:
    v = torch.zeros(STATE_BITS)
    i = 0
    for k in range(4):
        v[i] = (s["pc"] >> (3 - k)) & 1
        i += 1
    for r in range(4):
        for k in range(8):
            v[i] = (s["regs"][r] >> (7 - k)) & 1
            i += 1
    for k in range(4):
        v[i] = s["flags"][k]
        i += 1
    for k in range(4):
        v[i] = (s["sp"] >> (3 - k)) & 1
        i += 1
    v[i] = 1.0 if s["halted"] else 0.0
    i += 1
    for j in range(MEM_BYTES):
        for k in range(8):
            v[i] = (s["mem"][j] >> (7 - k)) & 1
            i += 1
    return v


def vec_to_state(v: torch.Tensor) -> dict:
    bits = [int(round(float(x))) for x in v.tolist()]
    i = 0
    pc = 0
    for k in range(4):
        pc = (pc << 1) | bits[i]
        i += 1
    regs = []
    for r in range(4):
        x = 0
        for k in range(8):
            x = (x << 1) | bits[i]
            i += 1
        regs.append(x)
    flags = bits[i:i + 4]
    i += 4
    sp = 0
    for k in range(4):
        sp = (sp << 1) | bits[i]
        i += 1
    halted = bool(bits[i])
    i += 1
    mem = []
    for j in range(MEM_BYTES):
        x = 0
        for k in range(8):
            x = (x << 1) | bits[i]
            i += 1
        mem.append(x)
    return {"pc": pc, "regs": regs, "flags": flags, "sp": sp,
            "halted": halted, "mem": mem}


# =============================================================================
# The step circuit: one full instruction of the registers-profile CPU
# =============================================================================

def build_step_net() -> Tuple[Net, List[str], List[str]]:
    net = Net()
    pc = [f"pc{i}" for i in range(4)]
    R = [[f"r{r}b{i}" for i in range(8)] for r in range(4)]
    fl = ["fZ", "fN", "fC", "fV"]
    sp = [f"sp{i}" for i in range(4)]
    halt = "halt"
    mem = [[f"m{j}b{i}" for i in range(8)] for j in range(MEM_BYTES)]

    # ---- fetch: bytes at PC .. PC+3 (decode lines reindexed, no adders) ----
    P = [net.DECODE(f"pdec{j}", pc, j) for j in range(MEM_BYTES)]
    byte = []
    for o in range(4):
        bo = []
        for bit in range(8):
            terms = [net.AND(f"f{o}a{j}b{bit}", [mem[j][bit], P[(j - o) % MEM_BYTES]])
                     for j in range(MEM_BYTES)]
            bo.append(net.OR(f"f{o}b{bit}", terms))
        byte.append(bo)

    # ---- field decode; opcode lines gated by NOT halt (total transition) ---
    nhalt = net.NOT("nhalt", halt)
    O = [net.DECODE(f"op{v:X}", byte[0][0:4], v) for v in range(16)]
    OG = [net.AND(f"og{v:X}", [O[v], nhalt]) for v in range(16)]
    rdl = [net.DECODE(f"rdl{r}", byte[0][4:6], r) for r in range(4)]
    rsl = [net.DECODE(f"rsl{r}", byte[0][6:8], r) for r in range(4)]
    CL = [net.DECODE(f"cond{c}", byte[1][5:8], c) for c in range(8)]

    # ---- register read ports ----------------------------------------------
    a = [net.OR(f"a{bit}", [net.AND(f"a{bit}r{r}", [R[r][bit], rdl[r]])
                            for r in range(4)]) for bit in range(8)]
    b = [net.OR(f"b{bit}", [net.AND(f"b{bit}r{r}", [R[r][bit], rsl[r]])
                            for r in range(4)]) for bit in range(8)]
    a_l = lambda k: a[7 - k]  # LSB-first views for the adder chains
    b_l = lambda k: b[7 - k]

    # ---- ALU: ADD ----------------------------------------------------------
    c = "#0"
    add_l = []
    for k in range(8):
        s_, c = net.FA(f"add.fa{k}", a_l(k), b_l(k), c)
        add_l.append(s_)
    add = [add_l[7 - i] for i in range(8)]
    add_c = c

    # ---- SUB (shared by CMP): a + ~b + 1 -----------------------------------
    nb_l = [net.NOT(f"sub.nb{k}", b_l(k)) for k in range(8)]
    c = "#1"
    sub_l = []
    for k in range(8):
        s_, c = net.FA(f"sub.fa{k}", a_l(k), nb_l[k], c)
        sub_l.append(s_)
    sub = [sub_l[7 - i] for i in range(8)]
    sub_c = c  # carry set = no borrow

    # ---- bitwise / shifts ---------------------------------------------------
    andr = [net.AND(f"land{bit}", [a[bit], b[bit]]) for bit in range(8)]
    orr = [net.OR(f"lor{bit}", [a[bit], b[bit]]) for bit in range(8)]
    xorr = [net.XOR(f"lxor{bit}", a[bit], b[bit]) for bit in range(8)]
    shl = [a[bit + 1] if bit < 7 else "#0" for bit in range(8)]
    shr = [a[bit - 1] if bit > 0 else "#0" for bit in range(8)]

    # ---- MUL: partial products + 7 chained shift-add adders ----------------
    pp = [[net.AND(f"pp.a{i}b{j}", [a[i], b[j]]) for j in range(8)] for i in range(8)]
    acc_l: List[str] = ["#0"] * 8
    acc_l[7] = pp[7][0]                      # row j=0 shifted by 7
    for j in range(1, 8):
        shift = 7 - j
        c = "#0"
        nxt = []
        for k in range(8):
            addend = pp[14 - k - j][j] if shift <= k <= min(7, 14 - j) else "#0"
            s_, c = net.FA(f"mul.s{j}.fa{k}", acc_l[k], addend, c)
            nxt.append(s_)
        acc_l = nxt
    mul = [acc_l[7 - i] for i in range(8)]

    # ---- DIV: 8 restoring stages (bit-cascade GE + subtract + mux) ---------
    def cascade_lt(prefix: str, x_m: List[str], y_m: List[str]) -> str:
        """Unsigned x < y via the verified bit-cascade (bit 0 = MSB)."""
        bit_lt, bit_eq = [], []
        for i in range(8):
            bit_lt.append(net.g(f"{prefix}.b{i}lt", [(x_m[i], -1), (y_m[i], 1)], -1)
                          if x_m[i] != "#0" else
                          (y_m[i] if y_m[i] != "#0" else "#0"))
            if x_m[i] == "#0":
                bit_eq.append(net.NOT(f"{prefix}.b{i}eq", y_m[i]))
            else:
                e_and = net.AND(f"{prefix}.b{i}ea", [x_m[i], y_m[i]])
                e_nor = net.NOR(f"{prefix}.b{i}en", [x_m[i], y_m[i]])
                bit_eq.append(net.OR(f"{prefix}.b{i}eq", [e_and, e_nor]))
        casc = [bit_lt[0]]
        for i in range(1, 8):
            pref = net.AND(f"{prefix}.ep{i}", bit_eq[:i])
            casc.append(net.AND(f"{prefix}.cl{i}", [pref, bit_lt[i]]))
        return net.OR(f"{prefix}.lt", casc)

    rem_l: List[str] = ["#0"] * 8
    q_m: List[str] = []
    for st in range(8):
        sh_l = [a[st]] + rem_l[0:7]          # (rem << 1) | a_bit; top bit provably 0
        sh_m = [sh_l[7 - i] for i in range(8)]
        lt = cascade_lt(f"div.s{st}.cmp", sh_m, b)
        ge = net.NOT(f"div.s{st}.ge", lt)
        c = "#1"
        dsub_l = []
        for k in range(8):
            s_, c = net.FA(f"div.s{st}.sub.fa{k}", sh_l[k], nb_l[k], c)
            dsub_l.append(s_)
        rem_l = [net.MUX(f"div.s{st}.mx{k}", ge, dsub_l[k], sh_l[k]) for k in range(8)]
        q_m.append(ge)
    div = q_m                                # quotient, MSB-first

    # ---- LOAD read port (address = low nibble of byte3) --------------------
    addr4 = byte[3][4:8]
    AD = [net.DECODE(f"adec{j}", addr4, j) for j in range(MEM_BYTES)]
    ld = [net.OR(f"ld{bit}", [net.AND(f"ld{bit}a{j}", [mem[j][bit], AD[j]])
                              for j in range(MEM_BYTES)]) for bit in range(8)]

    # ---- result mux over opcode lines ---------------------------------------
    srcs = [(0, add), (1, sub), (2, andr), (3, orr), (4, xorr),
            (5, shl), (6, shr), (7, mul), (8, div), (0xA, ld)]
    res = [net.OR(f"res{bit}", [net.AND(f"res{bit}o{v:X}", [src[bit], OG[v]])
                                for v, src in srcs]) for bit in range(8)]

    # ---- flags ---------------------------------------------------------------
    z_add = net.NOR("zadd", add)
    z_sub = net.NOR("zsub", sub)
    z_mul = net.NOR("zmul", mul)
    v_add = net.AND("vadd", [net.XOR("vadd.x1", a[0], add[0]),
                             net.XOR("vadd.x2", b[0], add[0])])
    v_sub = net.AND("vsub", [net.XOR("vsub.x1", a[0], b[0]),
                             net.XOR("vsub.x2", a[0], sub[0])])
    f_src = {  # per-source (Z, N, C, V); CMP shares the SUB datapath
        0: (z_add, add[0], add_c, v_add),
        1: (z_sub, sub[0], sub_c, v_sub),
        7: (z_mul, mul[0], "#0", "#0"),
        9: (z_sub, sub[0], sub_c, v_sub),
    }
    fw = net.OR("fw", [OG[0], OG[1], OG[7], OG[9]])
    fl_next = []
    for fi, fn in enumerate("ZNCV"):
        new = net.OR(f"fnew{fn}", [net.AND(f"fnew{fn}o{v:X}", [f_src[v][fi], OG[v]])
                                   for v in (0, 1, 7, 9)])
        fl_next.append(net.MUX(f"fmux{fn}", fw, new, fl[fi]))

    # ---- register writeback --------------------------------------------------
    wen = net.OR("wen", [OG[v] for v in (0, 1, 2, 3, 4, 5, 6, 7, 8, 0xA)])
    R_next = []
    for r in range(4):
        wsel = net.AND(f"wsel{r}", [wen, rdl[r]])
        nws = net.NOT(f"nws{r}", wsel)
        row = []
        for bit in range(8):
            t1 = net.AND(f"wb{r}b{bit}n", [res[bit], wsel])
            t0 = net.AND(f"wb{r}b{bit}o", [R[r][bit], nws])
            row.append(net.OR(f"wb{r}b{bit}", [t1, t0]))
        R_next.append(row)

    # ---- program counter ------------------------------------------------------
    pc_l = [pc[3 - k] for k in range(4)]     # LSB-first view
    p2_l = [pc_l[0], net.NOT("pc2.n1", pc_l[1]),
            net.XOR("pc2.x2", pc_l[2], pc_l[1]),
            net.XOR("pc2.x3", pc_l[3], net.AND("pc2.c3", [pc_l[2], pc_l[1]]))]
    p4_l = [pc_l[0], pc_l[1], net.NOT("pc4.n2", pc_l[2]),
            net.XOR("pc4.x3", pc_l[3], pc_l[2])]
    ext = net.OR("ext", [OG[v] for v in (0xA, 0xB, 0xC, 0xD, 0xE)])
    adv_l = [net.MUX(f"padv{k}", ext, p4_l[k], p2_l[k]) for k in range(4)]
    pc_adv = [adv_l[3 - i] for i in range(4)]

    nfl = [net.NOT(f"nf{fn}", fl[i]) for i, fn in enumerate("ZNCV")]
    cond_flag = [fl[0], nfl[0], fl[2], nfl[2], fl[1], nfl[1], fl[3], nfl[3]]
    condsel = net.OR("condsel", [net.AND(f"ct{cnd}", [CL[cnd], cond_flag[cnd]])
                                 for cnd in range(8)])
    take = net.OR("take", [OG[0xC], OG[0xE], net.AND("jcct", [OG[0xD], condsel])])
    keeppc = net.OR("keeppc", [halt, OG[0xF]])
    adv = net.AND("adv", [net.NOT("ntake", take), net.NOT("nkeep", keeppc)])
    pc_next = [net.OR(f"pcn{i}", [net.AND(f"pcn{i}t", [addr4[i], take]),
                                  net.AND(f"pcn{i}k", [pc[i], keeppc]),
                                  net.AND(f"pcn{i}a", [pc_adv[i], adv])])
               for i in range(4)]

    # ---- stack pointer (CALL: SP-2 = SP + 14) ---------------------------------
    sp_l = [sp[3 - k] for k in range(4)]
    c = "#0"
    s2_l = []
    for k, addend in enumerate(["#0", "#1", "#1", "#1"]):
        s_, c = net.FA(f"spm2.fa{k}", sp_l[k], addend, c)
        s2_l.append(s_)
    spm2 = [s2_l[3 - i] for i in range(4)]
    sp_next = [net.MUX(f"spn{i}", OG[0xE], spm2[i], sp[i]) for i in range(4)]

    # ---- halt (sticky) ---------------------------------------------------------
    halt_next = net.OR("haltn", [halt, OG[0xF]])

    # ---- memory next: STORE write, CALL pushes ret-hi (0) and ret-lo ----------
    SPD = [net.DECODE(f"spdec{j}", sp, j) for j in range(MEM_BYTES)]
    mem_next = []
    for j in range(MEM_BYTES):
        st_sel = net.AND(f"stsel{j}", [OG[0xB], AD[j]])
        ch_sel = net.AND(f"chsel{j}", [OG[0xE], SPD[(j + 1) % MEM_BYTES]])
        cl_sel = net.AND(f"clsel{j}", [OG[0xE], SPD[(j + 2) % MEM_BYTES]])
        keep = net.NOR(f"keep{j}", [st_sel, ch_sel, cl_sel])
        row = []
        for bit in range(8):
            terms = [net.AND(f"mn{j}b{bit}k", [mem[j][bit], keep]),
                     net.AND(f"mn{j}b{bit}s", [b[bit], st_sel])]
            if bit >= 4:  # ret-lo low nibble carries pc_adv; high nibble is 0
                terms.append(net.AND(f"mn{j}b{bit}c", [pc_adv[bit - 4], cl_sel]))
            row.append(net.OR(f"mn{j}b{bit}", terms))
        mem_next.append(row)

    outputs = (pc_next
               + [bitsig for row in R_next for bitsig in row]
               + fl_next + sp_next + [halt_next]
               + [bitsig for row in mem_next for bitsig in row])
    return net, state_names(), outputs


# =============================================================================
# Compiler: levelize ASAP, insert pass-throughs, emit dense ternary matrices
# =============================================================================

def compile_net(net: Net, inputs: List[str], outputs: List[str],
                ) -> Tuple[List[Tuple[torch.Tensor, torch.Tensor]], dict]:
    gates = net.gates
    level: Dict[str, int] = {s: 0 for s in inputs}

    consumers: Dict[str, List[str]] = {}
    indeg: Dict[str, int] = {}
    for gname, (ins, _) in gates.items():
        deps = [s for s, _ in ins if s in gates]
        indeg[gname] = len(deps)
        for s, _ in ins:
            consumers.setdefault(s, []).append(gname)
    order = [g for g, d in indeg.items() if d == 0]
    i = 0
    while i < len(order):
        for cns in consumers.get(order[i], []):
            indeg[cns] -= 1
            if indeg[cns] == 0:
                order.append(cns)
        i += 1
    if len(order) != len(gates):
        raise ValueError("cycle in step circuit")
    for gname in order:
        ins, _ = gates[gname]
        level[gname] = 1 + max((level[s] for s, _ in ins), default=0)

    kmax = max(level[gname] for gname in gates)
    K = kmax + 1  # final layer selects the next-state bits in order

    last_use: Dict[str, int] = {}
    for gname, (ins, _) in gates.items():
        for s, _ in ins:
            last_use[s] = max(last_use.get(s, 0), level[gname])
    for s in outputs:
        if s not in ("#0", "#1"):
            last_use[s] = K

    by_level: Dict[int, List[str]] = {}
    for gname in gates:
        by_level.setdefault(level[gname], []).append(gname)

    vec = list(inputs)                       # V_0
    layers: List[Tuple[torch.Tensor, torch.Tensor]] = []
    for lv in range(1, K + 1):
        if lv < K:
            new_gates = sorted(by_level.get(lv, []))
            passes = [s for s in vec if last_use.get(s, 0) > lv]
            rows = new_gates + passes
        else:
            rows = list(outputs)
        cols = {s: idx for idx, s in enumerate(vec)}
        W = torch.zeros(len(rows), len(vec), dtype=torch.int8)
        B = torch.zeros(len(rows), dtype=torch.int8)
        out_names = []
        for ri, s in enumerate(rows):
            if lv == K and s == "#0":
                B[ri] = -1
                out_names.append(s)
                continue
            if lv == K and s == "#1":
                B[ri] = 0
                out_names.append(s)
                continue
            if s in gates and level[s] == lv:
                ins, bias = gates[s]
                for src, w in ins:
                    W[ri, cols[src]] += w
                B[ri] = bias
            else:
                W[ri, cols[s]] = 1
                B[ri] = -1
            out_names.append(s)
        layers.append((W, B))
        vec = out_names

    if vec != outputs:
        raise AssertionError("final layer does not emit the state vector")

    info = {
        "layers": K,
        "gates": len(gates),
        "max_width": max(max(W.shape) for W, _ in layers),
        "total_weights": sum(W.numel() for W, _ in layers),
        "widths": [layers[0][0].shape[1]] + [W.shape[0] for W, _ in layers],
    }
    return layers, info


# =============================================================================
# Recurrent runtime
# =============================================================================

class MatrixMachine:
    """s' = H(Wk @ ... H(W1 @ s + b1) ... + bk); iterate until the halt bit.

    Digital mode thresholds at 0 (pre-activations are integers). Analog mode
    thresholds at -0.5 (the max-margin comparator point) and can inject
    Gaussian read noise per MVM and static ternary-conductance mismatch."""

    HALT_IDX = 4 + 32 + 4 + 4  # index of the halt bit in the state vector

    def __init__(self, layers, device="cpu"):
        self.device = device
        self.W = [W.to(device=device, dtype=torch.float32) for W, _ in layers]
        self.B = [B.to(device=device, dtype=torch.float32) for _, B in layers]

    @classmethod
    def from_file(cls, path=MODEL_PATH, device="cpu"):
        layers = []
        with safe_open(path, framework="pt") as f:
            n = 0
            while f"matrix.layer{n:03d}.weight" in f.keys():
                layers.append((f.get_tensor(f"matrix.layer{n:03d}.weight"),
                               f.get_tensor(f"matrix.layer{n:03d}.bias")))
                n += 1
        if not layers:
            raise FileNotFoundError(f"no matrix layers in {path}")
        return cls(layers, device=device)

    def perturbed(self, sigma_g: float, seed: int = 0) -> "MatrixMachine":
        """Static conductance mismatch: each nonzero ternary conductance gets a
        fixed Gaussian perturbation (fabrication variation, constant per run)."""
        gen = torch.Generator(device="cpu").manual_seed(seed)
        m = MatrixMachine.__new__(MatrixMachine)
        m.device = self.device
        m.W = []
        m.B = self.B
        for W in self.W:
            d = torch.randn(W.shape, generator=gen).to(self.device) * sigma_g
            m.W.append(W + d * (W != 0))
        return m

    def step(self, v: torch.Tensor, analog: bool = False,
             noise_sigma: float = 0.0, gen: Optional[torch.Generator] = None,
             ) -> torch.Tensor:
        thresh = -0.5 if analog else 0.0
        for W, b in zip(self.W, self.B):
            pre = v @ W.T + b
            if noise_sigma > 0.0:
                pre = pre + torch.randn(pre.shape, generator=gen,
                                        device=pre.device) * noise_sigma
            v = (pre >= thresh).float()
        return v

    def run(self, state: dict, max_cycles: int = 64, **kw) -> Tuple[dict, int]:
        v = state_to_vec(state).unsqueeze(0).to(self.device)
        n = 0
        while n < max_cycles and v[0, self.HALT_IDX] < 0.5:
            v = self.step(v, **kw)
            n += 1
        return vec_to_state(v[0].cpu()), n

    def min_margin(self, v: torch.Tensor) -> float:
        """Distance of every pre-activation from the analog threshold -0.5."""
        m = float("inf")
        for W, b in zip(self.W, self.B):
            pre = v @ W.T + b
            m = min(m, float((pre + 0.5).abs().min()))
            v = (pre >= 0).float()
        return m


# =============================================================================
# Pure-integer reference (mirrors GenericThresholdCPU semantics, 4-bit addr)
# =============================================================================

_JCC_FLAG = [0, 0, 2, 2, 1, 1, 3, 3]


def ref_step(s: dict) -> dict:
    if s["halted"]:
        return {k: (list(v) if isinstance(v, list) else v) for k, v in s.items()}
    s = {"pc": s["pc"], "regs": list(s["regs"]), "flags": list(s["flags"]),
         "sp": s["sp"], "halted": s["halted"], "mem": list(s["mem"])}
    M = MEM_BYTES - 1
    pc = s["pc"]
    ir = (s["mem"][pc] << 8) | s["mem"][(pc + 1) & M]
    op, rd, rs = (ir >> 12) & 0xF, (ir >> 10) & 3, (ir >> 8) & 3
    imm = ir & 0xFF
    next_pc = (pc + 2) & M
    addr = None
    if op in (0xA, 0xB, 0xC, 0xD, 0xE):
        al = s["mem"][(next_pc + 1) & M]
        addr = al & M
        next_pc = (next_pc + 2) & M
    a, b = s["regs"][rd], s["regs"][rs]
    result, carry, ovf, wr = a, 0, 0, True
    if op == 0x0:
        result = (a + b) & 0xFF
        carry = 1 if a + b > 0xFF else 0
        ovf = 1 if ((a ^ result) & (b ^ result)) & 0x80 else 0
    elif op == 0x1:
        result = (a - b) & 0xFF
        carry = 1 if a >= b else 0
        ovf = 1 if ((a ^ b) & (a ^ result)) & 0x80 else 0
    elif op == 0x2:
        result = a & b
    elif op == 0x3:
        result = a | b
    elif op == 0x4:
        result = a ^ b
    elif op == 0x5:
        result = (a << 1) & 0xFF
    elif op == 0x6:
        result = a >> 1
    elif op == 0x7:
        result = (a * b) & 0xFF
    elif op == 0x8:
        result = 0xFF if b == 0 else a // b
    elif op == 0x9:
        r2 = (a - b) & 0xFF
        carry = 1 if a >= b else 0
        ovf = 1 if ((a ^ b) & (a ^ r2)) & 0x80 else 0
        s["flags"] = [1 if r2 == 0 else 0, 1 if r2 & 0x80 else 0, carry, ovf]
        wr = False
    elif op == 0xA:
        result = s["mem"][addr]
    elif op == 0xB:
        s["mem"][addr] = b & 0xFF
        wr = False
    elif op == 0xC:
        s["pc"] = addr
        return s
    elif op == 0xD:
        cnd = imm & 7
        flag = s["flags"][_JCC_FLAG[cnd]]
        sel = flag if cnd % 2 == 0 else 1 - flag
        s["pc"] = addr if sel else next_pc
        return s
    elif op == 0xE:
        ret = next_pc
        sp2 = (s["sp"] - 1) & M
        s["mem"][sp2] = (ret >> 8) & 0xFF
        sp2 = (sp2 - 1) & M
        s["mem"][sp2] = ret & 0xFF
        s["sp"] = sp2
        s["pc"] = addr
        return s
    elif op == 0xF:
        s["halted"] = True
        return s
    if wr and op != 0x9:
        s["regs"][rd] = result & 0xFF
    if op in (0x0, 0x1, 0x7):
        s["flags"] = [1 if result == 0 else 0, 1 if result & 0x80 else 0,
                      carry, ovf]
    s["pc"] = next_pc
    return s


# =============================================================================
# Build + save
# =============================================================================

def build(save: bool = True):
    print("Assembling the step circuit from the verified cell library...")
    t0 = time.perf_counter()
    net, inputs, outputs = build_step_net()
    print(f"  {len(net.gates)} threshold gates")
    layers, info = compile_net(net, inputs, outputs)
    dt = time.perf_counter() - t0
    print(f"  compiled to {info['layers']} ternary matrices in {dt:.1f}s; "
          f"max width {info['max_width']}, "
          f"{info['total_weights']:,} dense weight entries")
    for W, B in layers:
        vals = set(torch.unique(W).tolist())
        assert vals <= {-1, 0, 1}, f"non-ternary matrix: {vals}"
        assert int(B.min()) >= -128 and int(B.max()) <= 127

    if save:
        tensors = {}
        for i, (W, B) in enumerate(layers):
            tensors[f"matrix.layer{i:03d}.weight"] = W
            tensors[f"matrix.layer{i:03d}.bias"] = B
        tensors["manifest.data_bits"] = torch.tensor([8.0])
        tensors["manifest.addr_bits"] = torch.tensor([float(ADDR_BITS)])
        tensors["manifest.memory_bytes"] = torch.tensor([float(MEM_BYTES)])
        tensors["manifest.registers"] = torch.tensor([4.0])
        tensors["manifest.layers"] = torch.tensor([float(info["layers"])])
        tensors["manifest.state_bits"] = torch.tensor([float(STATE_BITS)])
        tensors["manifest.version"] = torch.tensor([4.0])
        meta = {
            "machine": "matrix8",
            "weight_quantization": "ternary",
            "state_layout": json.dumps(state_names()),
            "analog": json.dumps({
                "comparator_threshold": -0.5,
                "guaranteed_margin": 0.5,
                "note": "pre-activations are integers; any total analog error "
                        "below 0.5 per pre-activation is provably bit-exact",
            }),
        }
        save_file(tensors, MODEL_PATH, metadata=meta)
        sz = os.path.getsize(MODEL_PATH) / (1024 * 1024)
        print(f"  saved {MODEL_PATH} ({sz:.1f} MB)")
    return layers, info


# =============================================================================
# Verification: matrix form == gate graph == integer reference, bit for bit
# =============================================================================

def _mk_state(mem=None, regs=(0, 0, 0, 0), flags=(0, 0, 0, 0), pc=0,
              sp=MEM_BYTES - 1, halted=False):
    m = [0] * MEM_BYTES
    if mem:
        for i, x in enumerate(mem):
            m[i] = x & 0xFF
    return {"pc": pc, "regs": list(regs), "flags": list(flags), "sp": sp,
            "halted": halted, "mem": m}


def _instr(op, rd=0, rs=0, imm=0):
    w = ((op & 0xF) << 12) | ((rd & 3) << 10) | ((rs & 3) << 8) | (imm & 0xFF)
    return [(w >> 8) & 0xFF, w & 0xFF]


def verify(mm: MatrixMachine, device: str) -> bool:
    ok_all = True

    # ---- 1. exhaustive ALU sweeps: all 65,536 operand pairs per opcode ------
    print("\n[1/4] Exhaustive ALU sweeps (65,536 operand pairs x 10 opcodes, "
          "batched through the matrices, vs integer reference)")
    aa = torch.arange(65536) % 256
    bb = torch.arange(65536) // 256
    for op, name in [(0, "ADD"), (1, "SUB"), (2, "AND"), (3, "OR"), (4, "XOR"),
                     (5, "SHL"), (6, "SHR"), (7, "MUL"), (8, "DIV"), (9, "CMP")]:
        for f_init in (0, 0xF):
            V = torch.zeros(65536, STATE_BITS)
            prog = _instr(op, rd=0, rs=1) + _instr(0xF)
            base = _mk_state(mem=prog, flags=[(f_init >> (3 - i)) & 1 for i in range(4)])
            V[:] = state_to_vec(base)
            for k in range(8):
                V[:, 4 + k] = ((aa >> (7 - k)) & 1).float()      # R0
                V[:, 12 + k] = ((bb >> (7 - k)) & 1).float()     # R1
            V = mm.step(V.to(device)).cpu()
            bad = 0
            probe = list(range(0, 65536, 4099)) + [0, 65535, 255, 256, 32768]
            for idx in probe:
                st = _mk_state(mem=prog, regs=(int(aa[idx]), int(bb[idx]), 0, 0),
                               flags=base["flags"])
                if vec_to_state(V[idx]) != ref_step(st):
                    bad += 1
            # full vectorized compare of the architectural result field
            r0 = sum(V[:, 4 + k].long() << (7 - k) for k in range(8))
            fZ, fN, fC, fV = (V[:, 36].long(), V[:, 37].long(),
                              V[:, 38].long(), V[:, 39].long())
            if op == 0:
                exp = (aa + bb) & 0xFF
            elif op in (1, 9):
                exp = (aa - bb) & 0xFF
            elif op == 2:
                exp = aa & bb
            elif op == 3:
                exp = aa | bb
            elif op == 4:
                exp = aa ^ bb
            elif op == 5:
                exp = (aa << 1) & 0xFF
            elif op == 6:
                exp = aa >> 1
            elif op == 7:
                exp = (aa * bb) & 0xFF
            else:
                exp = torch.where(bb == 0, torch.full_like(aa, 0xFF),
                                  aa // torch.clamp(bb, min=1))
            reg_exp = aa if op == 9 else exp
            n_bad = int((r0 != reg_exp).sum())
            if op in (0, 1, 7, 9):
                zx = (exp == 0).long()
                nx = ((exp & 0x80) > 0).long()
                if op == 0:
                    cx = ((aa + bb) > 0xFF).long()
                    vx = ((((aa ^ exp) & (bb ^ exp)) & 0x80) > 0).long()
                elif op in (1, 9):
                    cx = (aa >= bb).long()
                    vx = ((((aa ^ bb) & (aa ^ exp)) & 0x80) > 0).long()
                else:
                    cx = torch.zeros_like(aa)
                    vx = torch.zeros_like(aa)
                n_bad += int((fZ != zx).sum() + (fN != nx).sum()
                             + (fC != cx).sum() + (fV != vx).sum())
            else:
                want = torch.tensor([(f_init >> (3 - i)) & 1 for i in range(4)])
                n_bad += int((torch.stack([fZ, fN, fC, fV], 1)
                              != want.unsqueeze(0)).sum())
            status = "OK " if (n_bad == 0 and bad == 0) else "FAIL"
            if n_bad or bad:
                ok_all = False
            if f_init == 0:
                print(f"  {status} {name:4} 65,536/65,536 operand pairs exact"
                      + ("" if n_bad == 0 else f"  ({n_bad} mismatches)"))

    # ---- 2. exhaustive control-flow sweeps ----------------------------------
    print("\n[2/4] Exhaustive control-flow sweeps")
    cases = []
    for cnd in range(8):                     # all Jcc x all flag nibbles
        for f in range(16):
            prog = _instr(0xD, imm=cnd) + [0x00, 0x08] + _instr(0xF) + [0] * 2 + _instr(0xF)
            cases.append(_mk_state(mem=prog, flags=[(f >> (3 - i)) & 1 for i in range(4)]))
    for t in range(16):                      # all JMP targets
        prog = _instr(0xC) + [0x00, t]
        cases.append(_mk_state(mem=prog))
    for spv in range(16):                    # CALL: every SP value
        prog = _instr(0xE) + [0x00, 0x06] + [0, 0] + _instr(0xF)
        cases.append(_mk_state(mem=prog, sp=spv))
    for ad in range(16):                     # LOAD/STORE: every address
        prog = _instr(0xA, rd=2) + [0x00, ad]
        st = _mk_state(mem=prog)
        st["mem"][ad] |= 0xA5 if ad >= 4 else st["mem"][ad]
        cases.append(st)
        prog = _instr(0xB, rs=1) + [0x00, ad]
        cases.append(_mk_state(mem=prog, regs=(0, 0x5A, 0, 0)))
    V = torch.stack([state_to_vec(s) for s in cases]).to(device)
    V = mm.step(V).cpu()
    bad = sum(1 for i, s in enumerate(cases) if vec_to_state(V[i]) != ref_step(s))
    print(f"  {'OK ' if bad == 0 else 'FAIL'} {len(cases)} directed cases "
          f"(8 Jcc x 16 flag sets, 16 JMP targets, 16 CALL SPs, 32 LOAD/STORE) exact")
    ok_all &= bad == 0

    # ---- 3. random-state fuzz + halt fixed point -----------------------------
    print("\n[3/4] Random-state fuzz (single step vs integer reference)")
    gen = torch.Generator().manual_seed(0xC0FFEE)
    Vr = (torch.rand(2000, STATE_BITS, generator=gen) < 0.5).float()
    Vr[:, MatrixMachine.HALT_IDX] = 0.0
    out = mm.step(Vr.to(device)).cpu()
    bad = sum(1 for i in range(Vr.shape[0])
              if vec_to_state(out[i]) != ref_step(vec_to_state(Vr[i])))
    print(f"  {'OK ' if bad == 0 else 'FAIL'} 2,000 uniformly random full states exact")
    ok_all &= bad == 0
    Vh = (torch.rand(500, STATE_BITS, generator=gen) < 0.5).float()
    Vh[:, MatrixMachine.HALT_IDX] = 1.0
    outh = mm.step(Vh.to(device)).cpu()
    fixed = bool((outh == Vh.to(outh.device)).all())
    print(f"  {'OK ' if fixed else 'FAIL'} 500 halted states are exact fixed points")
    ok_all &= fixed

    # ---- 4. lockstep vs the gate-graph processor (shipped weights) ----------
    print("\n[4/4] Instruction lockstep vs GenericThresholdCPU on "
          "neural_computer8_registers (the gate graph)")
    from eval_all import GenericThresholdCPU
    T = {}
    with safe_open(GATE_MODEL_PATH, framework="pt") as f:
        for nm in f.keys():
            T[nm] = f.get_tensor(nm).float()
    gcpu = GenericThresholdCPU(T)
    progs = [
        ("alu_chain", _mk_state(mem=_instr(0, 0, 1) + _instr(7, 2, 0) + _instr(9, 2, 1)
                                + _instr(4, 1, 2) + _instr(0xF),
                                regs=(9, 5, 3, 0))),
        ("div_by_zero", _mk_state(mem=_instr(8, 0, 1) + _instr(0xF), regs=(77, 0, 0, 0))),
        ("load_store", _mk_state(mem=_instr(0xA, rd=3) + [0x00, 0x0E]
                                 + _instr(0xB, rs=3) + [0x00, 0x0F] + _instr(0xF)
                                 + [0] * 4 + [0xB7, 0x00],
                                 )),
        ("call", _mk_state(mem=_instr(0xE) + [0x00, 0x06] + _instr(0xF)
                           + _instr(0xF), sp=0xF)),
        ("jcc_taken", _mk_state(mem=_instr(0xD, imm=1) + [0x00, 0x06] + _instr(0xF)
                                + _instr(0xF), flags=(0, 0, 0, 0))),
        ("shift_mul", _mk_state(mem=_instr(5, 0) + _instr(6, 1) + _instr(7, 0, 1)
                                + _instr(0xF), regs=(0x81, 0x81, 0, 0))),
    ]
    for name, st0 in progs:
        sm = dict(st0)
        sg = {"pc": st0["pc"], "regs": list(st0["regs"]), "flags": list(st0["flags"]),
              "mem": list(st0["mem"]), "halted": False, "sp": st0["sp"]}
        good = True
        for _ in range(12):
            if sm["halted"]:
                break
            sm, _n = mm.run(sm, max_cycles=1)
            sg = gcpu.step(sg)
            match = (sm["pc"] == sg["pc"] and sm["regs"] == sg["regs"]
                     and sm["flags"] == sg["flags"] and sm["mem"] == sg["mem"]
                     and sm["halted"] == sg["halted"]
                     and sm["sp"] == sg.get("sp", MEM_BYTES - 1))
            if not match:
                good = False
                break
        good = good and sm["halted"]
        print(f"  {'OK ' if good else 'FAIL'} {name}: full-state lockstep, every instruction")
        ok_all &= good

    print("\nMATRIX == GATE GRAPH == REFERENCE:", "PASS" if ok_all else "FAIL")
    return ok_all


# =============================================================================
# Analog realization
# =============================================================================

def analog(mm: MatrixMachine, device: str) -> bool:
    ok = True
    print("\nAnalog crossbar simulation "
          "(comparator at -0.5; ternary conductances are exact by construction)")

    gen = torch.Generator().manual_seed(7)
    Vr = (torch.rand(512, STATE_BITS, generator=gen) < 0.5).float().to(device)
    margin = mm.min_margin(Vr)
    print(f"  measured minimum |pre-activation - (-0.5)| over 512 random states, "
          f"all layers: {margin:.3f} (guarantee: 0.5)")
    ok &= abs(margin - 0.5) < 1e-6

    prog = (_instr(0, 0, 1) + _instr(7, 2, 0) + _instr(9, 2, 1) + _instr(0xF))
    st0 = _mk_state(mem=prog, regs=(9, 5, 3, 0))
    ref_final, _ = mm.run(dict(st0), max_cycles=8)

    import math
    rows_per_step = sum(int(W.shape[0]) for W in mm.W)
    n_steps = 4
    print(f"  Gaussian read noise per MVM ({rows_per_step:,} comparator "
          f"decisions per instruction; flip iff |noise| >= 0.5):")
    for sigma in (0.05, 0.08, 0.10, 0.15, 0.25, 0.45):
        p_flip = math.erfc(0.5 / (sigma * math.sqrt(2.0)))
        exp_flips = p_flip * rows_per_step * n_steps
        exact = 0
        for t in range(20):
            g = torch.Generator(device=device).manual_seed(1000 + t)
            v = state_to_vec(st0).unsqueeze(0).to(device)
            for _ in range(8):
                if v[0, MatrixMachine.HALT_IDX] >= 0.5:
                    break
                v = mm.step(v, analog=True, noise_sigma=sigma, gen=g)
            exact += int(vec_to_state(v[0].cpu()) == ref_final)
        print(f"    sigma={sigma:.2f}: {exact}/20 runs bit-exact "
              f"(theory: {exp_flips:.2e} expected flips per run)")
        if exp_flips < 1e-3 and exact != 20:
            ok = False

    print("  Static ternary-conductance mismatch (device variation, "
          "10 fabricated instances each):")
    for sg in (0.02, 0.05, 0.10, 0.20):
        exact = 0
        for t in range(10):
            mmp = mm.perturbed(sg, seed=200 + t)
            v = state_to_vec(st0).unsqueeze(0).to(device)
            for _ in range(8):
                if v[0, MatrixMachine.HALT_IDX] >= 0.5:
                    break
                v = mmp.step(v, analog=True)
            exact += int(vec_to_state(v[0].cpu()) == ref_final)
        print(f"    sigma_G={sg:.2f}: {exact}/10 instances bit-exact")
        if sg <= 0.02 and exact != 10:
            ok = False

    print("ANALOG REALIZATION:", "PASS (bit-exact within the stated margins)"
          if ok else "FAIL")
    return ok


# =============================================================================
# CLI
# =============================================================================

def main() -> int:
    ap = argparse.ArgumentParser(description="neural_matrix8 builder/verifier")
    ap.add_argument("cmd", choices=["build", "verify", "analog", "all"])
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    args = ap.parse_args()

    rc = 0
    layers = None
    if args.cmd in ("build", "all"):
        layers, _ = build(save=True)
    if args.cmd in ("verify", "analog", "all"):
        if layers is not None:
            mm = MatrixMachine(layers, device=args.device)
        elif os.path.exists(MODEL_PATH):
            mm = MatrixMachine.from_file(MODEL_PATH, device=args.device)
        else:
            layers, _ = build(save=False)
            mm = MatrixMachine(layers, device=args.device)
        if args.cmd in ("verify", "all"):
            rc |= 0 if verify(mm, args.device) else 1
        if args.cmd in ("analog", "all"):
            rc |= 0 if analog(mm, args.device) else 1
    return rc


if __name__ == "__main__":
    sys.exit(main())