src/automata.py

"""The attractor, reversible, cellular-automaton and tile machines, with their builders."""
from __future__ import annotations

import json
import math
import os
import random
import sys
from collections import defaultdict
from typing import Dict, List, Optional, Tuple

import torch
from safetensors import safe_open
from safetensors.torch import save_file, load_file

from machines import Net, compile_net, MatrixMachine


class Circuit:
    """Wire allocator and energy accumulator."""

    def __init__(self) -> None:
        self.n = 0
        self.L: Dict[int, int] = defaultdict(int)
        self.Q: Dict[Tuple[int, int], int] = defaultdict(int)
        self.const = 0
        self.gates: List[Tuple[str, int, Tuple[int, ...]]] = []

    def wire(self) -> int:
        i = self.n
        self.n += 1
        return i

    def wires(self, k: int) -> List[int]:
        return [self.wire() for _ in range(k)]

    def _q(self, i: int, j: int, c: int) -> None:
        if i == j:
            self.L[i] += c
        else:
            self.Q[(min(i, j), max(i, j))] += c

    def AND(self, x: int, y: int) -> int:
        z = self.wire()
        self.L[z] += 3
        self._q(x, y, 1); self._q(x, z, -2); self._q(y, z, -2)
        self.gates.append(("AND", z, (x, y)))
        return z

    def OR(self, x: int, y: int) -> int:
        z = self.wire()
        self.L[x] += 1; self.L[y] += 1; self.L[z] += 1
        self._q(x, y, 1); self._q(x, z, -2); self._q(y, z, -2)
        self.gates.append(("OR", z, (x, y)))
        return z

    def NOT(self, x: int) -> int:
        z = self.wire()
        self.const += 1
        self.L[x] += -1; self.L[z] += -1
        self._q(x, z, 2)
        self.gates.append(("NOT", z, (x,)))
        return z

    def XOR(self, x: int, y: int) -> int:
        return self.OR(self.AND(x, self.NOT(y)), self.AND(self.NOT(x), y))

    def full_adder(self, x: int, y: int, cin: int) -> Tuple[int, int]:
        axy = self.XOR(x, y)
        s = self.XOR(axy, cin)
        cout = self.OR(self.AND(x, y), self.AND(cin, axy))
        return s, cout

    # ---- energy + couplings ------------------------------------------------
    def energy(self, s: List[int]) -> int:
        e = self.const
        for i, c in self.L.items():
            e += c * s[i]
        for (i, j), c in self.Q.items():
            e += c * s[i] * s[j]
        return e

    def neighbors(self) -> Dict[int, List[Tuple[int, int]]]:
        nbr: Dict[int, List[Tuple[int, int]]] = defaultdict(list)
        for (i, j), c in self.Q.items():
            nbr[i].append((j, c))
            nbr[j].append((i, c))
        return nbr

    # ---- relaxation modes --------------------------------------------------
    def forward_eval(self, clamp: Dict[int, int]) -> List[int]:
        """Exact forward relaxation from the clamped inputs to the energy-0 fixed point."""
        s = [0] * self.n
        for w, v in clamp.items():
            s[w] = v
        for op, z, ins in self.gates:
            if op == "AND":
                s[z] = s[ins[0]] & s[ins[1]]
            elif op == "OR":
                s[z] = s[ins[0]] | s[ins[1]]
            else:
                s[z] = 1 - s[ins[0]]
        return s

    def relax_energy(self, clamp: Dict[int, int], sweeps: int = 4000,
                     t0: float = 4.0, t1: float = 0.02, seed: int = 0
                     ) -> Tuple[List[int], bool]:
        """Anneal the full threshold network, tracking the lowest-energy state."""
        nbr = self.neighbors()
        rng = random.Random(seed)
        s = [rng.randint(0, 1) for _ in range(self.n)]
        for w, v in clamp.items():
            s[w] = v
        free = [i for i in range(self.n) if i not in clamp]
        best, best_e = list(s), self.energy(s)
        for step in range(sweeps):
            T = t0 * (t1 / t0) ** (step / max(1, sweeps - 1))
            for _ in range(len(free)):
                i = free[rng.randrange(len(free))]
                field = self.L[i] + sum(c * s[j] for j, c in nbr[i])
                dE = (1 - 2 * s[i]) * field
                if dE <= 0 or rng.random() < math.exp(-dE / T):
                    s[i] ^= 1
            e = self.energy(s)
            if e < best_e:
                best, best_e = list(s), e
                if best_e == 0:
                    return best, True
        return best, best_e == 0

    def solve(self, free_inputs: List[int], fixed: Dict[int, int],
              target: Dict[int, int], sweeps: int = 3000, restarts: int = 80,
              seed: int = 0) -> Optional[List[int]]:
        """Anneal chosen driver wires until the outputs match, the other wires slaved through the circuit."""
        rng = random.Random(seed)

        def mism(vals: Dict[int, int]) -> int:
            s = self.forward_eval({**fixed, **vals})
            return sum(1 for w, v in target.items() if s[w] != v)

        for _ in range(restarts):
            vals = {w: rng.randint(0, 1) for w in free_inputs}
            m = mism(vals)
            if m == 0:
                return self.forward_eval({**fixed, **vals})
            for step in range(sweeps):
                T = 2.0 * (0.02 / 2.0) ** (step / sweeps)
                w = free_inputs[rng.randrange(len(free_inputs))]
                vals[w] ^= 1
                m2 = mism(vals)
                if m2 <= m or rng.random() < math.exp(-(m2 - m) / T):
                    m = m2
                    if m == 0:
                        return self.forward_eval({**fixed, **vals})
                else:
                    vals[w] ^= 1
        return None


# ---------------------------------------------------------------------------
# Circuit builders
# ---------------------------------------------------------------------------
def adder(bits: int) -> Tuple[Circuit, dict]:
    c = Circuit()
    xs, ys = c.wires(bits), c.wires(bits)
    cin = c.wire()
    outs, carry = [], cin
    for k in range(bits):
        s, carry = c.full_adder(xs[k], ys[k], carry)
        outs.append(s)
    return c, {"xs": xs, "ys": ys, "cin": cin, "sum": outs + [carry]}


def multiplier(bits: int) -> Tuple[Circuit, dict]:
    c = Circuit()
    xs, ys = c.wires(bits), c.wires(bits)
    zero = c.wire()
    acc = [zero] * (2 * bits)
    for i in range(bits):
        carry = zero
        for j in range(bits):
            acc[i + j], carry = c.full_adder(acc[i + j], c.AND(xs[i], ys[j]), carry)
        acc[i + bits] = carry
    return c, {"xs": xs, "ys": ys, "zero": zero, "prod": acc}


_OPCODE = {"AND": 0, "OR": 1, "NOT": 2}

_OPNAME = {v: k for k, v in _OPCODE.items()}


def to_tensors(circ: Circuit, io: dict):
    """Serialize the coupling matrix (the program) and the gate list to tensors."""
    import torch
    qi = sorted(circ.Q)
    q_idx = torch.tensor(qi if qi else [], dtype=torch.long).reshape(-1, 2)
    q_val = torch.tensor([circ.Q[k] for k in qi], dtype=torch.long)
    li = sorted(circ.L)
    l_idx = torch.tensor(li, dtype=torch.long)
    l_val = torch.tensor([circ.L[i] for i in li], dtype=torch.long)
    g_op = torch.tensor([_OPCODE[op] for op, _, _ in circ.gates], dtype=torch.long)
    g_out = torch.tensor([o for _, o, _ in circ.gates], dtype=torch.long)
    g_in = torch.tensor([[ins[0], ins[1] if len(ins) > 1 else -1]
                         for _, _, ins in circ.gates], dtype=torch.long).reshape(-1, 2)
    t = {"Q_idx": q_idx, "Q_val": q_val, "L_idx": l_idx, "L_val": l_val,
         "gate_op": g_op, "gate_out": g_out, "gate_in": g_in}
    import json
    meta = {"machine": "attractor", "n": str(circ.n), "const": str(circ.const),
            "io": json.dumps({k: v for k, v in io.items()})}
    return t, meta


def from_tensors(t: dict, meta: dict) -> Tuple[Circuit, dict]:
    import json
    c = Circuit()
    c.n = int(meta["n"])
    c.const = int(meta["const"])
    for (i, j), v in zip(t["Q_idx"].tolist(), t["Q_val"].tolist()):
        c.Q[(i, j)] = v
    for i, v in zip(t["L_idx"].tolist(), t["L_val"].tolist()):
        c.L[i] = v
    for op, out, ins in zip(t["gate_op"].tolist(), t["gate_out"].tolist(), t["gate_in"].tolist()):
        c.gates.append((_OPNAME[op], out, tuple(x for x in ins if x >= 0)))
    return c, json.loads(meta["io"])


def cnf(clauses: List[List[int]], n_vars: int) -> Tuple[Circuit, dict]:
    """Compile a CNF formula."""
    c = Circuit()
    var = {v: c.wire() for v in range(1, n_vars + 1)}
    clause_ws = []
    for cl in clauses:
        lits = [var[abs(l)] if l > 0 else c.NOT(var[abs(l)]) for l in cl]
        acc = lits[0]
        for w in lits[1:]:
            acc = c.OR(acc, w)
        clause_ws.append(acc)
    sat = clause_ws[0]
    for w in clause_ws[1:]:
        sat = c.AND(sat, w)
    return c, {"vars": var, "sat": sat}


def test_forward():
    ok = True
    for bits in (4, 8):
        c, io = adder(bits)
        rng = random.Random(bits)
        bad = 0
        for _ in range(300):
            a, b = rng.randint(0, (1 << bits) - 1), rng.randint(0, (1 << bits) - 1)
            clamp = {io["cin"]: 0}
            for k in range(bits):
                clamp[io["xs"][k]] = (a >> k) & 1
                clamp[io["ys"][k]] = (b >> k) & 1
            s = c.forward_eval(clamp)
            got = sum(s[w] << k for k, w in enumerate(io["sum"]))
            if got != a + b or c.energy(s) != 0:
                bad += 1
        print(f"  forward adder {bits}-bit: {'OK' if bad == 0 else f'FAIL({bad})'}")
        ok &= bad == 0
    for bits in (3, 5):
        c, io = multiplier(bits)
        rng = random.Random(100 + bits)
        bad = 0
        for _ in range(300):
            a, b = rng.randint(0, (1 << bits) - 1), rng.randint(0, (1 << bits) - 1)
            clamp = {io["zero"]: 0}
            for k in range(bits):
                clamp[io["xs"][k]] = (a >> k) & 1
                clamp[io["ys"][k]] = (b >> k) & 1
            s = c.forward_eval(clamp)
            got = sum(s[w] << k for k, w in enumerate(io["prod"]))
            if got != a * b or c.energy(s) != 0:
                bad += 1
        print(f"  forward multiplier {bits}-bit: {'OK' if bad == 0 else f'FAIL({bad})'}")
        ok &= bad == 0
    return ok


def test_energy_relax():
    """The canonical form: anneal the whole network (no propagation shortcut)."""
    c, io = adder(4)
    rng = random.Random(3)
    bad = 0
    for _ in range(20):
        a, b = rng.randint(0, 15), rng.randint(0, 15)
        clamp = {io["cin"]: 0}
        for k in range(4):
            clamp[io["xs"][k]] = (a >> k) & 1
            clamp[io["ys"][k]] = (b >> k) & 1
        conv = False
        for attempt in range(4):                 # annealers restart
            s, conv = c.relax_energy(clamp, sweeps=6000, seed=rng.randint(0, 1 << 30))
            got = sum(s[w] << k for k, w in enumerate(io["sum"]))
            if conv and got == a + b:
                break
        if not conv:
            bad += 1
    print(f"  whole-network energy relaxation (4-bit adder, 20 cases): "
          f"{'OK' if bad == 0 else f'reached ground state in {20 - bad}/20'}")
    return bad == 0


def test_factor():
    ok = True
    for bits, targets in ((4, [15, 35, 143]), (5, [21, 55, 91])):
        c, io = multiplier(bits)
        for N in targets:
            target = {io["prod"][k]: (N >> k) & 1 for k in range(2 * bits)}
            s = c.solve(io["xs"] + io["ys"], {io["zero"]: 0}, target, seed=N)
            if s is None:
                print(f"  factor {N}: not found")
                ok = False
                continue
            a = sum(s[io["xs"][k]] << k for k in range(bits))
            b = sum(s[io["ys"][k]] << k for k in range(bits))
            good = a * b == N and 1 < a < N and 1 < b < N
            print(f"  factor {N} ({bits}x{bits}): {a} x {b} {'OK' if a * b == N else 'WRONG'}")
            ok &= a * b == N
    return ok


def test_sat():
    # (x1 | x2 | ~x3) & (~x1 | x3) & (x2 | x3) & (~x2 | ~x3), a satisfiable 3-SAT.
    clauses = [[1, 2, -3], [-1, 3], [2, 3], [-2, -3]]
    c, io = cnf(clauses, 3)
    s = c.solve(list(io["vars"].values()), {}, {io["sat"]: 1}, seed=1)
    if s is None:
        print("  SAT solve: no model found")
        return False
    assign = {v: s[w] for v, w in io["vars"].items()}
    sat = all(any((assign[abs(l)] == 1) if l > 0 else (assign[abs(l)] == 0) for l in cl)
              for cl in clauses)
    print(f"  SAT solve: model {assign} {'satisfies' if sat else 'FAILS'} the formula")
    return sat

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


ATTRACTOR_OUT = os.path.join(ROOT, "variants", "neural_attractor.safetensors")

BITS = 8


def main_build_attractor() -> int:
    c, io = multiplier(BITS)
    tensors, meta = to_tensors(c, io)
    save_file(tensors, ATTRACTOR_OUT, metadata=meta)
    size = os.path.getsize(ATTRACTOR_OUT)
    print(f"Built {os.path.relpath(ATTRACTOR_OUT, ROOT)}: {BITS}x{BITS} multiplier as couplings")
    print(f"  wires={c.n}  Q entries={len(c.Q)}  L entries={len(c.L)}  size={size/1024:.1f} KB")

    # round-trip
    t = load_file(ATTRACTOR_OUT)
    with safe_open(ATTRACTOR_OUT, framework="pt") as f:
        m = f.metadata()
    c2, io2 = from_tensors(t, m)
    xs, ys, zero, prod = io2["xs"], io2["ys"], io2["zero"], io2["prod"]

    rng = random.Random(0)
    bad = 0
    for _ in range(200):
        a, b = rng.randint(0, 255), rng.randint(0, 255)
        clamp = {zero: 0}
        for k in range(BITS):
            clamp[xs[k]] = (a >> k) & 1
            clamp[ys[k]] = (b >> k) & 1
        s = c2.forward_eval(clamp)
        got = sum(s[w] << k for k, w in enumerate(prod))
        if got != a * b or c2.energy(s) != 0:
            bad += 1
    print(f"  round-trip forward multiply (200 cases): {'OK' if bad == 0 else f'FAIL({bad})'}")

    # backward: factor a small product through the loaded couplings
    N = 35
    target = {prod[k]: (N >> k) & 1 for k in range(2 * BITS)}
    s = c2.solve(xs + ys, {zero: 0}, target, sweeps=2500, restarts=120, seed=N)
    if s is not None:
        fa = sum(s[xs[k]] << k for k in range(BITS))
        fb = sum(s[ys[k]] << k for k in range(BITS))
        print(f"  round-trip backward factor {N}: {fa} x {fb} {'OK' if fa * fb == N else 'WRONG'}")
        ok = fa * fb == N
    else:
        print(f"  round-trip backward factor {N}: not found")
        ok = False
    return 0 if (bad == 0 and ok) else 1

# --- threshold-gate truth: the target updates are computed by Heaviside gates ---
def H(x: int) -> int:
    return 1 if x >= 0 else 0


def g_and(a: int, b: int) -> int:
    return H(a + b - 2)          # fires iff a=b=1


def g_xor(a: int, b: int) -> int:
    # XOR as AND(OR, NAND): OR=H(a+b-1), NAND=H(1-a-b), out=AND(OR,NAND)
    return g_and(H(a + b - 1), H(1 - a - b))


# --- reversible primitives, in place on a bit register (a Python list) ---
def NOT(reg: List[int], t: int) -> None:
    reg[t] = 1 - reg[t]


def CNOT(reg: List[int], c: int, t: int) -> None:
    reg[t] = g_xor(reg[t], reg[c])


def TOFF(reg: List[int], a: int, b: int, t: int) -> None:
    reg[t] = g_xor(reg[t], g_and(reg[a], reg[b]))


def FRED(reg: List[int], c: int, x: int, y: int) -> None:
    # controlled swap of x,y on control c
    if reg[c]:
        reg[x], reg[y] = reg[y], reg[x]


# --- Cuccaro ripple-carry adder: b += a in place, one carry ancilla ---
# MAJ(c,b,a): b^=a ; c^=a ; a ^= b&c        UMA(c,b,a): a ^= b&c ; c^=a ; b^=c
def _maj(reg, c, b, a):
    CNOT(reg, a, b)
    CNOT(reg, a, c)
    TOFF(reg, b, c, a)


def _uma(reg, c, b, a):
    TOFF(reg, b, c, a)
    CNOT(reg, a, c)
    CNOT(reg, c, b)


# Each primitive below is its own inverse, so the inverse of a gate sequence is
# the reversed sequence. Building the adder as an op list gives subtraction for
# free (run it backward).
def _maj_ops(c, b, a):
    return [(CNOT, a, b), (CNOT, a, c), (TOFF, b, c, a)]


def _uma_ops(c, b, a):
    return [(TOFF, b, c, a), (CNOT, a, c), (CNOT, c, b)]


def _adder_ops(a_bits, b_bits, carry, cout=None):
    n = len(a_bits)
    ops = _maj_ops(carry, b_bits[0], a_bits[0])
    for i in range(1, n):
        ops += _maj_ops(a_bits[i - 1], b_bits[i], a_bits[i])
    if cout is not None:
        ops.append((CNOT, a_bits[n - 1], cout))
    for i in range(n - 1, 0, -1):
        ops += _uma_ops(a_bits[i - 1], b_bits[i], a_bits[i])
    ops += _uma_ops(carry, b_bits[0], a_bits[0])
    return ops


def _apply(reg, ops, inverse=False):
    for gate, *args in (reversed(ops) if inverse else ops):
        gate(reg, *args)


def add_into(reg, a_bits, b_bits, carry, cout=None):
    """b <- (a + b) mod 2^n (LSB first); a and carry (=0) restored."""
    _apply(reg, _adder_ops(a_bits, b_bits, carry, cout))


def sub_into(reg, a_bits, b_bits, carry, cout=None):
    """b <- (b - a) mod 2^n: the adder run backward."""
    _apply(reg, _adder_ops(a_bits, b_bits, carry, cout), inverse=True)


def xor_into(reg, a_bits, b_bits):
    """b <- b XOR a, bitwise (self-inverse)."""
    for a, b in zip(a_bits, b_bits):
        CNOT(reg, a, b)


def incr(reg, b_bits, one_bits, carry):
    """b <- b + 1 mod 2^n."""
    add_into(reg, one_bits, b_bits, carry)


def neg_into(reg, b_bits, one_bits, carry):
    """b <- (-b) mod 2^n via two's complement (~b then +1). Self-inverse."""
    for t in b_bits:
        NOT(reg, t)
    add_into(reg, one_bits, b_bits, carry)


def rot_left(reg, b_bits, k=1):
    """Rotate the word left by k (a permutation of bit positions; reversible)."""
    n = len(b_bits)
    k %= n
    vals = [reg[b_bits[i]] for i in range(n)]
    for i in range(n):
        reg[b_bits[(i + k) % n]] = vals[i]


def and_into(reg, a_bits, b_bits, t_bits):
    """t <- t XOR (a AND b), bitwise (word-level Toffoli). Self-inverse."""
    for a, b, t in zip(a_bits, b_bits, t_bits):
        TOFF(reg, a, b, t)


# --- verification helpers ---
def is_permutation(fn, nbits: int) -> bool:
    """Check fn: {0,1}^nbits -> {0,1}^nbits is a bijection (exhaustive)."""
    seen = set()
    for x in range(1 << nbits):
        reg = [(x >> k) & 1 for k in range(nbits)]
        fn(reg)
        y = sum(b << k for k, b in enumerate(reg))
        seen.add(y)
    return len(seen) == (1 << nbits)


def _test_primitives():
    ok = True
    ok &= is_permutation(lambda r: NOT(r, 0), 1)
    ok &= is_permutation(lambda r: CNOT(r, 0, 1), 2)
    ok &= is_permutation(lambda r: TOFF(r, 0, 1, 2), 3)
    ok &= is_permutation(lambda r: FRED(r, 0, 1, 2), 3)
    print(f"  primitives bijective (NOT/CNOT/TOFF/FRED): {'OK' if ok else 'FAIL'}")
    return ok


def _test_adder(width=4):
    # layout: a[0..w-1], b[0..w-1], carry
    a_bits = list(range(width))
    b_bits = list(range(width, 2 * width))
    carry = 2 * width
    n = 2 * width + 1
    ok_perm = is_permutation(lambda r: add_into(r, a_bits, b_bits, carry), n)
    bad = 0
    mask = (1 << width) - 1
    for a in range(1 << width):
        for b in range(1 << width):
            reg = [0] * n
            for k in range(width):
                reg[a_bits[k]] = (a >> k) & 1
                reg[b_bits[k]] = (b >> k) & 1
            add_into(reg, a_bits, b_bits, carry)
            got_a = sum(reg[a_bits[k]] << k for k in range(width))
            got_b = sum(reg[b_bits[k]] << k for k in range(width))
            if got_a != a or got_b != ((a + b) & mask) or reg[carry] != 0:
                bad += 1
    print(f"  Cuccaro adder {width}-bit: bijection={'OK' if ok_perm else 'FAIL'}  "
          f"b<-a+b, a & carry restored={'OK' if bad == 0 else f'FAIL({bad})'}")
    return ok_perm and bad == 0


def bennett(reg, a_b, b_b, c_b, scr_b, out_b, carry):
    """Bennett compute-copy-uncompute for the irreversible f(a,b,c)=(a+b) XOR c."""
    ops = []
    # compute f into scratch (scratch starts 0): scratch += a; scratch += b; scratch ^= c
    ops += _adder_ops(a_b, scr_b, carry)
    ops += _adder_ops(b_b, scr_b, carry)
    ops += [(CNOT, c, s) for c, s in zip(c_b, scr_b)]
    # copy scratch -> out
    ops += [(CNOT, s, o) for s, o in zip(scr_b, out_b)]
    # uncompute scratch (reverse of the compute prefix)
    n_comp = len(_adder_ops(a_b, scr_b, carry)) * 2 + len(c_b)
    ops += [(g, *args) for g, *args in reversed(ops[:n_comp])]
    _apply(reg, ops)
    return ops


def _test_bennett(width=4):
    a_b = list(range(width))
    b_b = list(range(width, 2 * width))
    c_b = list(range(2 * width, 3 * width))
    scr_b = list(range(3 * width, 4 * width))
    out_b = list(range(4 * width, 5 * width))
    carry = 5 * width
    n = 5 * width + 1
    mask = (1 << width) - 1
    bad = 0
    for a in range(1 << width):
        for b in range(1 << width):
            for c in range(1 << width):
                r = [0] * n
                for k in range(width):
                    r[a_b[k]] = (a >> k) & 1
                    r[b_b[k]] = (b >> k) & 1
                    r[c_b[k]] = (c >> k) & 1
                bennett(r, a_b, b_b, c_b, scr_b, out_b, carry)
                ra = sum(r[a_b[k]] << k for k in range(width))
                rb = sum(r[b_b[k]] << k for k in range(width))
                rc = sum(r[c_b[k]] << k for k in range(width))
                rs = sum(r[scr_b[k]] << k for k in range(width))
                ro = sum(r[out_b[k]] << k for k in range(width))
                if (ra, rb, rc, rs, ro, r[carry]) != (a, b, c, 0, ((a + b) & mask) ^ c, 0):
                    bad += 1
    print(f"  Bennett (a,b,c,0,0)->(a,b,c,0,(a+b)^c), scratch cleaned "
          f"[{width}-bit]: {'OK' if bad == 0 else f'FAIL({bad})'}")
    return bad == 0


def _test_alu(width=4):
    mask = (1 << width) - 1
    a_b = list(range(width))
    b_b = list(range(width, 2 * width))
    one_b = list(range(2 * width, 3 * width))
    t_b = list(range(3 * width, 4 * width))
    carry = 4 * width
    n = 4 * width + 1

    def fresh(a=0, b=0, t=0, one=False):
        r = [0] * n
        for k in range(width):
            r[a_b[k]] = (a >> k) & 1
            r[b_b[k]] = (b >> k) & 1
            r[t_b[k]] = (t >> k) & 1
        if one:
            r[one_b[0]] = 1
        return r

    def rd(r, bits):
        return sum(r[bits[k]] << k for k in range(width))

    results = {}
    # subtract
    bad = 0
    for a in range(1 << width):
        for b in range(1 << width):
            r = fresh(a, b)
            sub_into(r, a_b, b_b, carry)
            if rd(r, b_b) != ((b - a) & mask) or rd(r, a_b) != a or r[carry] != 0:
                bad += 1
    results["sub b-=a"] = bad == 0 and is_permutation(lambda r: sub_into(r, a_b, b_b, carry), n)
    # xor
    bad = 0
    for a in range(1 << width):
        for b in range(1 << width):
            r = fresh(a, b)
            xor_into(r, a_b, b_b)
            if rd(r, b_b) != (a ^ b) or rd(r, a_b) != a:
                bad += 1
    results["xor b^=a"] = bad == 0
    # negate
    bad = 0
    for b in range(1 << width):
        r = fresh(b=b, one=True)
        neg_into(r, b_b, one_b, carry)
        if rd(r, b_b) != ((-b) & mask) or rd(r, one_b) != 1 or r[carry] != 0:
            bad += 1
    results["neg b=-b"] = bad == 0
    # increment
    bad = 0
    for b in range(1 << width):
        r = fresh(b=b, one=True)
        incr(r, b_b, one_b, carry)
        if rd(r, b_b) != ((b + 1) & mask):
            bad += 1
    results["incr b+=1"] = bad == 0
    # rotate
    bad = 0
    for b in range(1 << width):
        r = fresh(b=b)
        rot_left(r, b_b, 1)
        exp = ((b << 1) | (b >> (width - 1))) & mask
        if rd(r, b_b) != exp:
            bad += 1
    results["rot_left"] = bad == 0
    # and-into (Toffoli word)
    bad = 0
    for a in range(1 << width):
        for b in range(1 << width):
            for t in range(1 << width):
                r = fresh(a, b, t)
                and_into(r, a_b, b_b, t_b)
                if rd(r, t_b) != (t ^ (a & b)):
                    bad += 1
    results["and t^=a&b"] = bad == 0
    ok = all(results.values())
    print("  reversible ALU " + f"{width}-bit: " +
          "  ".join(f"{k}={'OK' if v else 'FAIL'}" for k, v in results.items()))
    return ok


class RCPU:
    def __init__(self, program: List[tuple], k_regs=4, width=8, mem_words=16):
        self.prog = program
        self.K = k_regs
        self.W = width
        self.M = mem_words
        self.mask = (1 << width) - 1
        self.L = len(program)

    def new_state(self, regs=None, mem=None) -> dict:
        return {"R": list(regs) + [0] * (self.K - len(regs)) if regs else [0] * self.K,
                "PC": 0, "BR": 1, "MEM": list(mem) + [0] * (self.M - len(mem)) if mem else [0] * self.M}

    # ---- reversible instruction effects (forward and inverse) ----
    def _data(self, s, I, inverse: bool):
        R, MEM, m = s["R"], s["MEM"], self.mask
        op = I[0]
        if op == "ADD":
            d, r = I[1], I[2]
            R[d] = (R[d] - R[r]) & m if inverse else (R[d] + R[r]) & m
        elif op == "SUB":
            d, r = I[1], I[2]
            R[d] = (R[d] + R[r]) & m if inverse else (R[d] - R[r]) & m
        elif op == "ADDI":
            d, k = I[1], I[2]
            R[d] = (R[d] - k) & m if inverse else (R[d] + k) & m
        elif op == "XOR":
            R[I[1]] ^= R[I[2]]
        elif op == "XORI":
            R[I[1]] ^= (I[2] & m)
        elif op == "NEG":
            R[I[1]] = (-R[I[1]]) & m
        elif op == "TOFF":
            R[I[1]] ^= (R[I[2]] & R[I[3]])
        elif op == "ROL":
            k = (-I[2] if inverse else I[2]) % self.W
            R[I[1]] = ((R[I[1]] << k) | (R[I[1]] >> (self.W - k))) & m if k else R[I[1]]
        elif op == "EXCH":
            d, r = I[1], I[2]
            a = R[r] % self.M
            R[d], MEM[a] = MEM[a], R[d]
        # BRA/BEZ/HALT have no data effect

    def _toggle(self, s, I):
        """Reversible control: toggle BR for a taken branch (self-inverse)."""
        op = I[0]
        if op == "BRA":
            s["BR"] ^= (I[1] ^ 1)
        elif op == "BEZ":
            if s["R"][I[1]] == 0:
                s["BR"] ^= (I[2] ^ 1)

    # ---- single-step transition and its inverse ----
    def step(self, s):
        I = self.prog[s["PC"]]
        self._data(s, I, inverse=False)
        self._toggle(s, I)
        s["PC"] = (s["PC"] + s["BR"]) % self.L

    def step_back(self, s):
        s["PC"] = (s["PC"] - s["BR"]) % self.L
        I = self.prog[s["PC"]]
        self._toggle(s, I)               # self-inverse: restores BR
        self._data(s, I, inverse=True)

    def run(self, s, steps):
        for _ in range(steps):
            self.step(s)
        return s

    def run_back(self, s, steps):
        for _ in range(steps):
            self.step_back(s)
        return s


def _clone(s):
    return {"R": list(s["R"]), "PC": s["PC"], "BR": s["BR"], "MEM": list(s["MEM"])}


def _eq(a, b):
    return a["R"] == b["R"] and a["PC"] == b["PC"] and a["BR"] == b["BR"] and a["MEM"] == b["MEM"]


def test_straight_line():
    prog = [("ADD", 1, 0), ("XOR", 1, 0), ("NEG", 1), ("ADDI", 1, 7),
            ("TOFF", 2, 0, 1), ("ROL", 0, 1)]
    m = RCPU(prog, width=8)
    ok = True
    for a in (5, 0, 255, 100):
        for b in (3, 1, 200):
            s0 = m.new_state([a, b, 0, 0])
            s = _clone(s0)
            m.run(s, len(prog))
            fwd = _clone(s)
            m.run_back(s, len(prog))
            ok &= _eq(s, s0)               # round-trip recovers the exact input
    print(f"  straight-line round-trip (dir -1 recovers input): {'OK' if ok else 'FAIL'}")
    return ok


def test_bijection():
    """The single-step transition is a bijection: step_back inverts step on every instruction."""
    import itertools
    W = 4
    prog = [
        ("ADD", 1, 0), ("SUB", 1, 0), ("ADDI", 2, 3), ("XOR", 1, 2),
        ("XORI", 0, 5), ("NEG", 3), ("TOFF", 3, 0, 1), ("ROL", 0, 1),
        ("EXCH", 2, 3), ("BRA", 2), ("BEZ", 1, 3), ("BEZ", 0, -2),
    ]
    m = RCPU(prog, k_regs=4, width=W, mem_words=4)
    bad = 0
    checked = 0
    rng = __import__("random").Random(0)
    for pc in range(len(prog)):
        for br in (1, 2, 3, -2, -1):
            for _ in range(60):
                s0 = {"R": [rng.randint(0, (1 << W) - 1) for _ in range(4)],
                      "PC": pc, "BR": br,
                      "MEM": [rng.randint(0, (1 << W) - 1) for _ in range(4)]}
                s = _clone(s0)
                m.step(s)
                m.step_back(s)
                if not _eq(s, s0):
                    bad += 1
                # and the other composition order
                s = _clone(s0)
                m.step_back(s)
                m.step(s)
                if not _eq(s, s0):
                    bad += 1
                checked += 2
    print(f"  step_back o step = id over {checked} (state, instruction) cases: "
          f"{'OK' if bad == 0 else f'FAIL({bad})'}")
    return bad == 0

WIDTH = 8

MASK = (1 << WIDTH) - 1


def _regs_written(stmts) -> set:
    w = set()
    for st in stmts:
        op = st[0]
        if op in ("ADD", "SUB", "ADDI", "XOR", "NEG"):
            w.add(st[1])
        elif op == "SWAP":
            w.add(st[1]); w.add(st[2])
        elif op == "FOR":
            w |= _regs_written(st[2])
        elif op == "IF":
            w |= _regs_written(st[2]) | _regs_written(st[3])
    return w


def invert(stmts: List[tuple]) -> List[tuple]:
    out = []
    for st in reversed(stmts):
        op = st[0]
        if op == "ADD":
            out.append(("SUB", st[1], st[2]))
        elif op == "SUB":
            out.append(("ADD", st[1], st[2]))
        elif op == "ADDI":
            out.append(("ADDI", st[1], -st[2]))
        elif op in ("XOR", "NEG", "SWAP"):
            out.append(st)
        elif op == "FOR":
            out.append(("FOR", st[1], invert(st[2])))
        elif op == "IF":
            # Janus: reverse swaps predicate and exit assertion, inverts branches
            out.append(("IF", st[4], invert(st[2]), invert(st[3]), st[1]))
    return out


def prog_run(stmts: List[tuple], s: Dict[str, int]) -> Dict[str, int]:
    for st in stmts:
        op = st[0]
        if op == "ADD":
            s[st[1]] = (s[st[1]] + s[st[2]]) & MASK
        elif op == "SUB":
            s[st[1]] = (s[st[1]] - s[st[2]]) & MASK
        elif op == "ADDI":
            s[st[1]] = (s[st[1]] + st[2]) & MASK
        elif op == "XOR":
            s[st[1]] ^= s[st[2]]
        elif op == "NEG":
            s[st[1]] = (-s[st[1]]) & MASK
        elif op == "SWAP":
            s[st[1]], s[st[2]] = s[st[2]], s[st[1]]
        elif op == "FOR":
            cnt, body = st[1], st[2]
            if cnt in _regs_written(body):
                raise ValueError("FOR count register must be read-only (irreversible otherwise)")
            for _ in range(s[cnt]):
                prog_run(body, s)
        elif op == "IF":
            pred, then, els, exit_assert = st[1], st[2], st[3], st[4]
            if pred(s):
                prog_run(then, s)
                assert exit_assert(s), "exit assertion violated (not reversible)"
            else:
                prog_run(els, s)
                assert not exit_assert(s), "exit assertion violated (not reversible)"
    return s


# --- demonstration programs ---
MULTIPLY = [("FOR", "b", [("ADD", "acc", "a")])]          # acc += a, b times; a,b preserved

FIB = [("FOR", "n", [("ADD", "a", "b"), ("SWAP", "a", "b")])]   # (a,b)->(b,a+b), n times


def _test():
    ok = True

    # reversible multiply: acc = a*b, inputs preserved; inverse clears acc
    bad = 0
    for a in range(16):
        for b in range(16):
            s = {"a": a, "b": b, "acc": 0}
            prog_run(MULTIPLY, s)
            if s["acc"] != (a * b) & MASK or s["a"] != a or s["b"] != b:
                bad += 1
            prog_run(invert(MULTIPLY), s)                       # run backward
            if s != {"a": a, "b": b, "acc": 0}:
                bad += 1
    print(f"  reversible multiply acc=a*b, inputs preserved, inverse clears acc: "
          f"{'OK' if bad == 0 else f'FAIL({bad})'}")
    ok &= bad == 0

    # reversible Fibonacci: n steps forward, inverse recovers the seed
    bad = 0
    for n in range(12):
        s = {"a": 0, "b": 1, "n": n}
        prog_run(FIB, s)
        # forward value check against a plain reference
        ra, rb = 0, 1
        for _ in range(n):
            ra, rb = rb, (ra + rb) & MASK
        if (s["a"], s["b"]) != (ra, rb):
            bad += 1
        prog_run(invert(FIB), s)
        if s != {"a": 0, "b": 1, "n": n}:
            bad += 1
    print(f"  reversible Fibonacci n steps, inverse recovers seed: "
          f"{'OK' if bad == 0 else f'FAIL({bad})'}")
    ok &= bad == 0

    # a reversible conditional (Janus IF): swap when the operands differ. The
    # exit assertion a != b is true exactly when the then-branch ran (swapping
    # distinct values keeps them distinct; equal values are skipped), so the
    # reverse picks the right branch.
    prog = [("IF", lambda s: s["a"] != s["b"], [("SWAP", "a", "b")], [],
             lambda s: s["a"] != s["b"])]
    bad = 0
    for a in range(16):
        for b in range(16):
            s = {"a": a, "b": b}
            prog_run(prog, s)
            if sorted([s["a"], s["b"]]) != sorted([a, b]):
                bad += 1
            prog_run(invert(prog), s)
            if s != {"a": a, "b": b}:
                bad += 1
    print(f"  reversible conditional (Janus IF with exit assertion): "
          f"{'OK' if bad == 0 else f'FAIL({bad})'}")
    ok &= bad == 0
    return ok


REVERSIBLE_OUT = os.path.join(ROOT, "variants", "neural_reversible.safetensors")

# The two reversible primitives, keyed to their gate functions.
_CODE = {CNOT: 0, TOFF: 1}

_FN = {0: CNOT, 1: TOFF}

# Heaviside threshold gates that implement the target updates: a CNOT target is
# XOR(t,c); a Toffoli target is XOR(t, AND(a,b)). AND/OR/NAND are single gates.
_SUBSTRATE = {
    "and_w": torch.tensor([1, 1]), "and_b": torch.tensor(-2),
    "or_w": torch.tensor([1, 1]), "or_b": torch.tensor(-1),
    "nand_w": torch.tensor([-1, -1]), "nand_b": torch.tensor(1),
}


def encode(ops):
    codes, args = [], []
    for gate, *a in ops:
        codes.append(_CODE[gate])
        args.append((a + [-1, -1, -1])[:3])
    return torch.tensor(codes, dtype=torch.long), torch.tensor(args, dtype=torch.long)


def main_build_reversible() -> int:
    a_bits = list(range(WIDTH))
    b_bits = list(range(WIDTH, 2 * WIDTH))
    carry = 2 * WIDTH
    n = 2 * WIDTH + 1
    ops = _adder_ops(a_bits, b_bits, carry)
    codes, args = encode(ops)

    tensors = {"gate_code": codes, "gate_args": args, **_SUBSTRATE}
    import json
    meta = {"machine": "reversible", "width": str(WIDTH), "n_wires": str(n),
            "a_bits": json.dumps(a_bits), "b_bits": json.dumps(b_bits),
            "carry": str(carry), "circuit": "in-place reversible adder b<-a+b"}
    save_file(tensors, REVERSIBLE_OUT, metadata=meta)
    print(f"Built {os.path.relpath(REVERSIBLE_OUT, ROOT)}: reversible {WIDTH}-bit adder")
    print(f"  gates={len(ops)}  wires={n}  size={os.path.getsize(REVERSIBLE_OUT)} bytes")

    # round-trip: reconstruct the op list from the file and run it
    t = load_file(REVERSIBLE_OUT)
    with safe_open(REVERSIBLE_OUT, framework="pt") as f:
        m = f.metadata()
    W = int(m["width"]); ab = json.loads(m["a_bits"]); bb = json.loads(m["b_bits"])
    cy = int(m["carry"]); nn = int(m["n_wires"])
    loaded = []
    for code, a in zip(t["gate_code"].tolist(), t["gate_args"].tolist()):
        loaded.append((_FN[code], *[x for x in a if x >= 0]))

    def run(reg):
        for gate, *a in loaded:
            gate(reg, *a)

    mask = (1 << W) - 1
    bad = 0
    rng = random.Random(0)
    for _ in range(400):
        a = rng.randint(0, mask); b = rng.randint(0, mask)
        reg = [0] * nn
        for k in range(W):
            reg[ab[k]] = (a >> k) & 1
            reg[bb[k]] = (b >> k) & 1
        run(reg)
        got_b = sum(reg[bb[k]] << k for k in range(W))
        got_a = sum(reg[ab[k]] << k for k in range(W))
        if got_b != ((a + b) & mask) or got_a != a or reg[cy] != 0:
            bad += 1
    print(f"  round-trip b<-a+b, addend & carry restored (400 cases): "
          f"{'OK' if bad == 0 else f'FAIL({bad})'}")
    # bijection is exhaustive-checkable at width 4 on the same construction
    a4 = list(range(4)); b4 = list(range(4, 8)); c4 = 8
    ops4 = _adder_ops(a4, b4, c4)
    perm = is_permutation(lambda r: _apply(r, ops4), 9)
    print(f"  loaded circuit is a bijection (4-bit): {'OK' if perm else 'FAIL'}")
    return 0 if (bad == 0 and perm) else 1


def adder_net(width: int):
    """The in-place Cuccaro adder as a feedforward ternary netlist computing its permutation."""
    a_bits = list(range(width))
    b_bits = list(range(width, 2 * width))
    carry = 2 * width
    n = 2 * width + 1
    ops = _adder_ops(a_bits, b_bits, carry)
    net = Net()
    inputs = [f"in{i}" for i in range(n)]
    cur = list(inputs)
    for k, (gate, *args) in enumerate(ops):
        if gate is CNOT:
            c, t = args
            cur[t] = net.XOR(f"c{k}", cur[t], cur[c])          # t ^= c
        else:                                                   # Toffoli: t ^= a&b
            a, b, t = args
            tmp = net.AND(f"t{k}a", [cur[a], cur[b]])
            cur[t] = net.XOR(f"t{k}x", cur[t], tmp)
    return net, inputs, list(cur), n, a_bits, b_bits, carry


def _refs(n, a_bits, b_bits, carry):
    out = []
    for x in range(1 << n):
        reg = [(x >> i) & 1 for i in range(n)]
        _apply(reg, _adder_ops(a_bits, b_bits, carry))
        out.append(sum(reg[i] << i for i in range(n)))
    return torch.tensor(out)


def _outputs(mm, states, n, **step_kw):
    v = mm.step(states.clone(), **step_kw)
    bits = (v >= 0.5).to(torch.int64)
    weights = torch.tensor([1 << i for i in range(n)])
    return (bits * weights).sum(dim=1)


def analog_sweep(mm, states, refs, n):
    """The permutation must survive analog imperfection."""
    print("  read-noise sweep (bit-exact = all 512 outputs still match, 4 trials):")
    for sigma in (0.05, 0.10, 0.15, 0.20):
        ok = all((_outputs(mm, states, n, analog=True, noise_sigma=sigma,
                           gen=torch.Generator().manual_seed(s)) == refs).all()
                 for s in range(4))
        print(f"    sigma={sigma:.2f}: {'bit-exact' if ok else 'errors appear'}")
    print("  conductance-mismatch sweep (fixed per-device weight perturbation):")
    for sg in (0.05, 0.10, 0.15):
        ok = (_outputs(mm.perturbed(sg, seed=0), states, n, analog=True) == refs).all()
        print(f"    sigma_G={sg:.2f}: {'bit-exact' if ok else 'errors appear'}")


def main_reversible_matrix() -> int:
    width = 4
    net, inputs, outputs, n, a_bits, b_bits, carry = adder_net(width)
    layers, info = compile_net(net, inputs, outputs)
    mm = MatrixMachine(layers)

    states = torch.stack([torch.tensor([float((x >> i) & 1) for i in range(n)])
                          for x in range(1 << n)])
    refs = _refs(n, a_bits, b_bits, carry)
    got = _outputs(mm, states, n)
    bad = int((got != refs).sum())
    perm = len(set(got.tolist())) == (1 << n)
    margin = mm.min_margin(states[:256])

    print(f"reversible {width}-bit adder as a ternary matrix stack")
    print(f"  layers={info['layers']}  gates={info['gates']}  "
          f"max_width={info['max_width']}  total_weights={info['total_weights']}")
    print(f"  every weight ternary: {'OK' if all(((W == -1) | (W == 0) | (W == 1)).all() for W, _ in layers) else 'FAIL'}")
    print(f"  matches the gate circuit over all {1 << n} inputs: {'OK' if bad == 0 else f'FAIL({bad})'}")
    print(f"  composed transition is a permutation of the state space: {'OK' if perm else 'FAIL'}")
    print(f"  analog noise margin, all layers: {margin:.3f} (guarantee 0.5)")
    analog_sweep(mm, states, refs, n)
    ok = bad == 0 and perm and abs(margin - 0.5) < 1e-6
    print("PASS" if ok else "FAIL")
    return 0 if ok else 1

Block = Tuple[int, int, int, int]


def rule(b: Block) -> Block:
    tl, tr, bl, br = b
    if b == (1, 0, 0, 1):
        return (0, 1, 1, 0)          # diagonal pair -> other diagonal (deflect)
    if b == (0, 1, 1, 0):
        return (1, 0, 0, 1)
    return (br, bl, tr, tl)          # otherwise rotate 180 degrees


def is_bijection() -> bool:
    outs = {rule(tuple((s >> k) & 1 for k in (3, 2, 1, 0))) for s in range(16)}
    return len(outs) == 16


def self_inverse() -> bool:
    return all(rule(rule(tuple((s >> k) & 1 for k in (3, 2, 1, 0))))
               == tuple((s >> k) & 1 for k in (3, 2, 1, 0)) for s in range(16))


def step(grid: List[List[int]], phase: int) -> List[List[int]]:
    """One Margolus update; phase 0 aligns blocks at even coordinates, phase 1 offsets them by (1,1)."""
    H, W = len(grid), len(grid[0])
    out = [row[:] for row in grid]
    o = phase
    for r0 in range(o, o + H, 2):
        for c0 in range(o, o + W, 2):
            r, r1 = r0 % H, (r0 + 1) % H
            c, c1 = c0 % W, (c0 + 1) % W
            nb = rule((grid[r][c], grid[r][c1], grid[r1][c], grid[r1][c1]))
            out[r][c], out[r][c1], out[r1][c], out[r1][c1] = nb
    return out


def ca_run(grid: List[List[int]], nsteps: int, start_phase: int = 0) -> List[List[int]]:
    g = grid
    for n in range(nsteps):
        g = step(g, (start_phase + n) & 1)
    return g


def run_back(grid: List[List[int]], nsteps: int, start_phase: int = 0) -> List[List[int]]:
    """Undo `run`: replay the phase sequence in reverse; the rule is self-inverse."""
    phases = [(start_phase + n) & 1 for n in range(nsteps)]
    g = grid
    for p in reversed(phases):
        g = step(g, p)
    return g


# --- the block rule as Heaviside threshold gates ---
# rule(s) = rotate180(s) XOR is_diag(s) on every cell: rotation fixes diagonal
# pairs, and flipping all four cells of a rotated diagonal pair sends it to the
# other diagonal. is_diag detects the two diagonal-pair states.
def _H(x):
    return 1 if x >= 0 else 0


def _and(*xs):
    return _H(sum(xs) - len(xs))


def _or(*xs):
    return _H(sum(xs) - 1)


def _xor(a, b):
    return _and(_or(a, b), _H(1 - a - b))       # OR AND NAND, the family's XOR


def gate_rule(b: Block) -> Block:
    tl, tr, bl, br = b
    d = _or(_and(tl, 1 - tr, 1 - bl, br), _and(1 - tl, tr, bl, 1 - br))   # is_diag
    return (_xor(br, d), _xor(bl, d), _xor(tr, d), _xor(tl, d))


def _test_gates():
    ok = all(gate_rule(tuple((s >> k) & 1 for k in (3, 2, 1, 0)))
             == rule(tuple((s >> k) & 1 for k in (3, 2, 1, 0))) for s in range(16))
    print(f"  block rule as Heaviside threshold gates matches over 16 states: "
          f"{'OK' if ok else 'FAIL'}")
    return ok


# --- tests ---
def _rand_grid(H, W, seed):
    import random
    rng = random.Random(seed)
    return [[rng.randint(0, 1) for _ in range(W)] for _ in range(H)]


def _ball_positions(g):
    return {(r, c) for r, row in enumerate(g) for c, v in enumerate(row) if v}


def _test_rule():
    print(f"  block rule is a bijection of 16 states: {'OK' if is_bijection() else 'FAIL'}")
    print(f"  block rule is self-inverse: {'OK' if self_inverse() else 'FAIL'}")
    return is_bijection() and self_inverse()


def _test_reversibility():
    bad = 0
    for seed in range(20):
        g = _rand_grid(8, 8, seed)
        fwd = ca_run(g, 25, start_phase=0)
        back = run_back(fwd, 25, start_phase=0)
        if back != g:
            bad += 1
    # particle count is conserved (the rule permutes cells within each block)
    g = _rand_grid(8, 8, 99)
    conserved = sum(sum(r) for r in g) == sum(sum(r) for r in ca_run(g, 40))
    print(f"  lattice reversible (run then reverse recovers grid, 20 grids): "
          f"{'OK' if bad == 0 else f'FAIL({bad})'}")
    print(f"  particle number conserved: {'OK' if conserved else 'FAIL'}")
    return bad == 0 and conserved


def _test_ballistic():
    # a single particle travels in a straight diagonal line
    H = W = 16
    g = [[0] * W for _ in range(H)]
    g[2][2] = 1
    positions = [next(iter(_ball_positions(g)))]
    gg = g
    for n in range(8):
        gg = step(gg, n & 1)
        p = _ball_positions(gg)
        positions.append(next(iter(p)) if len(p) == 1 else None)
    ok = all(p is not None for p in positions)
    steady = ok and all(positions[i + 1] == (positions[i][0] + 1, positions[i][1] + 1)
                        for i in range(len(positions) - 1))
    print(f"  single particle stays a single particle: {'OK' if ok else 'FAIL'}")
    print(f"  and moves ballistically on the diagonal, +(1,1) per step: "
          f"{'OK' if steady else 'FAIL'}  trace={positions[:5]}")
    return ok and steady


def interaction_gate(a: int, b: int) -> dict:
    """One billiard-ball collision as a reversible interaction gate."""
    H = W = 12
    g = [[0] * W for _ in range(H)]
    if a:
        g[2][2] = 1
    if b:
        g[7][7] = 1
    g = ca_run(g, 4)
    return {"A_and_B": g[3][6], "A_and_notB": g[6][6], "notA_and_B": g[3][3]}


def and_gate(a: int, b: int) -> int:
    return interaction_gate(a, b)["A_and_B"]


def _test_gate():
    ok = True
    for a in (0, 1):
        for b in (0, 1):
            o = interaction_gate(a, b)
            ok &= (o["A_and_B"] == (a & b) and o["A_and_notB"] == (a & (1 - b))
                   and o["notA_and_B"] == ((1 - a) & b))
    print(f"  billiard-ball interaction gate (A&B, A&~B, ~A&B) over all 4 inputs: "
          f"{'OK' if ok else 'FAIL'}")
    return ok


def and3(a: int, b: int, c: int) -> int:
    """Two composed collisions."""
    H = W = 16
    g = [[0] * W for _ in range(H)]
    if a:
        g[2][2] = 1
    if b:
        g[7][7] = 1
    if c:
        g[0][9] = 1
    return ca_run(g, 4)[3][5]


def _test_compose():
    ok = all(and3(a, b, c) == (a & b & c)
             for a in (0, 1) for b in (0, 1) for c in (0, 1))
    print(f"  composed 3-input AND (two chained collisions) over all 8 inputs: "
          f"{'OK' if ok else 'FAIL'}")
    return ok


def _test_collision():
    # Two particles interact (the joint evolution differs from independent
    # motion) and the collision stays reversible: the physics that logic needs.
    H = W = 12
    interacted = False
    revok = True
    # converging pairs: an SE-mover (even,even) meets an NW-mover (odd,odd) on a
    # shared diagonal, forming the 1001/0110 diagonal pair the rule deflects.
    for a, b in [((2, 2), (7, 7)), ((3, 3), (8, 8)), ((2, 8), (7, 3)),
                 ((4, 4), (9, 9)), ((2, 2), (9, 9))]:
        g = [[0] * W for _ in range(H)]
        g[a[0]][a[1]] = 1
        g[b[0]][b[1]] = 1
        ga = [[0] * W for _ in range(H)]
        ga[a[0]][a[1]] = 1
        gb = [[0] * W for _ in range(H)]
        gb[b[0]][b[1]] = 1
        joint = g
        for n in range(12):
            joint = step(joint, n & 1)
            ga = step(ga, n & 1)
            gb = step(gb, n & 1)
            free = _ball_positions(ga) | _ball_positions(gb)
            if _ball_positions(joint) != free:
                interacted = True
        if run_back(ca_run(g, 12), 12) != g:
            revok = False
    print(f"  two-particle collisions interact (joint != independent motion): "
          f"{'OK' if interacted else 'FAIL'}")
    print(f"  collisions remain reversible: {'OK' if revok else 'FAIL'}")
    return interacted and revok

CA_OUT = os.path.join(ROOT, "variants", "neural_ca.safetensors")


def block_net():
    net = Net()
    ntl, ntr = net.NOT("ntl", "tl"), net.NOT("ntr", "tr")
    nbl, nbr = net.NOT("nbl", "bl"), net.NOT("nbr", "br")
    d1 = net.AND("d1", ["tl", ntr, nbl, "br"])
    d2 = net.AND("d2", [ntl, "tr", "bl", nbr])
    dg = net.OR("dg", [d1, d2])                       # is_diag
    TL = net.XOR("oTL", "br", dg)
    TR = net.XOR("oTR", "bl", dg)
    BL = net.XOR("oBL", "tr", dg)
    BR = net.XOR("oBR", "tl", dg)
    return net, ["tl", "tr", "bl", "br"], [TL, TR, BL, BR]


def matrix_step(mm, grid, phase):
    """One Margolus update driven entirely by the matrix tile."""
    H, W = len(grid), len(grid[0])
    out = [row[:] for row in grid]
    for r0 in range(phase, phase + H, 2):
        for c0 in range(phase, phase + W, 2):
            r, r1, c, c1 = r0 % H, (r0 + 1) % H, c0 % W, (c0 + 1) % W
            v = torch.tensor([[float(grid[r][c]), float(grid[r][c1]),
                               float(grid[r1][c]), float(grid[r1][c1])]])
            o = mm.step(v)[0]
            out[r][c], out[r][c1], out[r1][c], out[r1][c1] = (int(o[0]), int(o[1]),
                                                              int(o[2]), int(o[3]))
    return out


def main_build_ca() -> int:
    net, inp, outp = block_net()
    layers, info = compile_net(net, inp, outp)
    mm = MatrixMachine(layers)

    seen, bad, vecs = set(), 0, []
    for s in range(16):
        b = tuple((s >> k) & 1 for k in (3, 2, 1, 0))     # tl,tr,bl,br
        v = torch.tensor([[float(x) for x in b]])
        got = tuple(int(x) for x in mm.step(v)[0])
        if got != rule(b):
            bad += 1
        seen.add(got)
        vecs.append(v[0])
    perm = len(seen) == 16
    margin = mm.min_margin(torch.stack(vecs))

    tensors = {}
    for k, (W, B) in enumerate(layers):
        tensors[f"matrix.layer{k:03d}.weight"] = W.to(torch.int8)
        tensors[f"matrix.layer{k:03d}.bias"] = B.to(torch.int8)
    meta = {"machine": "ca",
            "rule": "Margolus reversible: rotate 180 except diagonal pair swaps (BBM class)",
            "inputs": "tl,tr,bl,br", "outputs": "TL,TR,BL,BR", "layers": str(info["layers"])}
    save_file(tensors, CA_OUT, metadata=meta)
    print(f"Built {os.path.relpath(CA_OUT, ROOT)}: reversible CA block rule as a ternary matrix tile")
    print(f"  layers={info['layers']}  gates={info['gates']}  size={os.path.getsize(CA_OUT)} bytes")
    print(f"  every weight ternary: "
          f"{'OK' if all(((W == -1) | (W == 0) | (W == 1)).all() for W, _ in layers) else 'FAIL'}")
    print(f"  tile matches the block rule over all 16 states: {'OK' if bad == 0 else f'FAIL({bad})'}")
    print(f"  tile is a permutation (16 distinct outputs): {'OK' if perm else 'FAIL'}")
    print(f"  analog noise margin: {margin:.3f} (guarantee 0.5)")

    # analog robustness: the tile must reproduce the rule under read noise and
    # static conductance mismatch, as neural_matrix8 measures
    states = torch.stack(vecs)
    ref = [rule(tuple((s >> k) & 1 for k in (3, 2, 1, 0))) for s in range(16)]

    def outs(machine, **kw):
        v = machine.step(states, **kw)
        return [tuple(int(v[i][j]) for j in range(4)) for i in range(v.shape[0])]

    print("  read-noise sweep (exact = tile matches the rule on all 16 states):")
    for sigma in (0.05, 0.10, 0.15, 0.20):
        okn = all(outs(mm, analog=True, noise_sigma=sigma,
                       gen=torch.Generator().manual_seed(s)) == ref for s in range(4))
        print(f"    sigma={sigma:.2f}: {'exact' if okn else 'errors appear'}")
    print("  conductance-mismatch sweep:")
    for sg in (0.05, 0.10, 0.15):
        okg = outs(mm.perturbed(sg, seed=0), analog=True) == ref
        print(f"    sigma_G={sg:.2f}: {'exact' if okg else 'errors appear'}")

    # the loaded tile, applied to every block, is one whole-lattice CA step
    t = load_file(CA_OUT)
    n = 0
    lyr = []
    while f"matrix.layer{n:03d}.weight" in t:
        lyr.append((t[f"matrix.layer{n:03d}.weight"], t[f"matrix.layer{n:03d}.bias"]))
        n += 1
    mm2 = MatrixMachine(lyr)
    g = _rand_grid(8, 8, 3)
    gmat = matrix_step(mm2, g, 0)
    gref = step(g, 0)
    print(f"  loaded tile drives a full lattice step (matches ca.step): "
          f"{'OK' if gmat == gref else 'FAIL'}")

    ok = bad == 0 and perm and abs(margin - 0.5) < 1e-6 and gmat == gref
    print("PASS" if ok else "FAIL")
    return 0 if ok else 1

# side -> (dx, dy, my_side, neighbor_side)
_SIDES = [(0, 1, "N", "S"), (1, 0, "E", "W"), (0, -1, "S", "N"), (-1, 0, "W", "E")]


class Tile:
    __slots__ = ("N", "E", "S", "W", "name")

    def __init__(self, N="", E="", S="", W="", name=""):
        self.N, self.E, self.S, self.W, self.name = N, E, S, W, name

    def glue(self, side):
        return getattr(self, side)


def bind_strength(A: Dict[Tuple[int, int], Tile], x: int, y: int, t: Tile,
                  strength: Dict[str, int]) -> int:
    """Summed strength of t's glues that match the abutting neighbor glues."""
    s = 0
    for dx, dy, side, opp in _SIDES:
        nb = A.get((x + dx, y + dy))
        if nb is None:
            continue
        g = t.glue(side)
        if g and g == nb.glue(opp):
            s += strength.get(g, 1)
    return s


def binds(A, x, y, t, tau, strength) -> bool:
    """The threshold-gate binding decision: H(sum strength*match - tau)."""
    return bind_strength(A, x, y, t, strength) >= tau


def grow(tileset: List[Tile], seed: Dict[Tuple[int, int], Tile], tau: int,
         strength: Dict[str, int], bounds: Tuple[int, int, int, int],
         max_tiles: int = 100000) -> Tuple[Dict[Tuple[int, int], Tile], bool]:
    """Directed growth from a seed."""
    x0, y0, x1, y1 = bounds
    A = dict(seed)
    deterministic = True
    changed = True
    while changed and len(A) < max_tiles:
        changed = False
        frontier = set()
        for (x, y) in list(A):
            for dx, dy, _, _ in _SIDES:
                p = (x + dx, y + dy)
                if p not in A and x0 <= p[0] <= x1 and y0 <= p[1] <= y1:
                    frontier.add(p)
        for (x, y) in frontier:
            binders = [t for t in tileset if binds(A, x, y, t, tau, strength)]
            if len(binders) == 1:
                A[(x, y)] = binders[0]
                changed = True
            elif len(binders) > 1:
                deterministic = False
    return A, deterministic


# ---------------------------------------------------------------------------
# XOR / Sierpinski tile set: value(x,y) = value(x-1,y) XOR value(x,y-1)
# ---------------------------------------------------------------------------
def rule2_tileset(fn) -> List[Tile]:
    """Rule tiles for value(x,y) = fn(W, S): four cooperatively binding tiles emitting fn north and east."""
    ts = []
    for s in (0, 1):
        for w in (0, 1):
            v = fn(w, s)
            ts.append(Tile(N=f"v{v}", E=f"v{v}", S=f"v{s}", W=f"v{w}",
                           name=f"R w{w} s{s} -> {v}"))
    return ts


def sierpinski_tileset() -> List[Tile]:
    return rule2_tileset(lambda w, s: w ^ s)


def _row_col_seed(bottom: List[int], left: List[int]):
    """Seed the bottom row and the left column with fixed value tiles."""
    seed = {}
    for x, b in enumerate(bottom):
        seed[(x, 0)] = Tile(N=f"v{b}", E="", S="", W="", name=f"seedB{x}={b}")
    for y, l in enumerate(left):
        if y == 0:
            continue
        seed[(0, y)] = Tile(N="", E=f"v{l}", S="", W="", name=f"seedL{y}={l}")
    return seed


def _test_binding_gate():
    """The binding decision is exactly the Heaviside threshold gate."""
    strength = {"v0": 1, "v1": 1}
    ts = sierpinski_tileset()
    A = {(1, 0): Tile(N="v1"), (0, 1): Tile(E="v0")}
    bad = 0
    for t in ts:
        for x, y in [(1, 1)]:
            w = sum(strength.get(t.glue(side), 1)
                    for dx, dy, side, opp in _SIDES
                    if A.get((x + dx, y + dy)) and t.glue(side)
                    and t.glue(side) == A[(x + dx, y + dy)].glue(opp))
            gate = 1 if (w - 2) >= 0 else 0                # H(sum*match - tau)
            if gate != int(binds(A, x, y, t, 2, strength)):
                bad += 1
    print(f"  binding decision == Heaviside gate H(sum-tau): {'OK' if bad == 0 else 'FAIL'}")
    return bad == 0


def _test_rule2(fn, name, n=24):
    strength = {"v0": 1, "v1": 1}
    bottom = [1 if x == 0 else 0 for x in range(n)]
    left = [1 if y == 0 else 0 for y in range(n)]
    seed = _row_col_seed(bottom, left)
    A, det = grow(rule2_tileset(fn), seed, 2, strength, (0, 0, n - 1, n - 1))

    def val(x, y):
        t = A.get((x, y))
        return None if t is None else (1 if t.N == "v1" else 0)

    ref = {(x, 0): bottom[x] for x in range(n)}
    ref.update({(0, y): left[y] for y in range(n)})
    for y in range(1, n):
        for x in range(1, n):
            ref[(x, y)] = fn(ref[(x - 1, y)], ref[(x, y - 1)])

    filled = bad = 0
    for y in range(1, n):
        for x in range(1, n):
            v = val(x, y)
            if v is not None:
                filled += 1
                bad += v != ref[(x, y)]
    tag = "OK" if (det and bad == 0 and filled > 0) else "FAIL"
    print(f"  rule-tile CA fn={name:3s}: directed={det}  placed={filled}  "
          f"every tile = fn(W,S) {tag}")
    return det and bad == 0 and filled > 0


# ---------------------------------------------------------------------------
# Binary counter: each row is the row below plus one. LSB is the right column;
# carry propagates west by cooperative binding (S = bit below, E = carry in).
# ---------------------------------------------------------------------------
def counter_tileset() -> List[Tile]:
    ts = []
    for b in (0, 1):
        for c in (0, 1):
            ts.append(Tile(N=f"b{b ^ c}", E=f"c{c}", S=f"b{b}", W=f"c{b & c}",
                           name=f"C b{b} c{c} -> b{b ^ c} carry{b & c}"))
    ts.append(Tile(N="edge", E="", S="edge", W="c1", name="edge(+1 injector)"))
    return ts


def counter_seed(n: int):
    """Bottom row (y=0) all zero, plus the right-edge +1 injector column base."""
    seed = {}
    for x in range(n):
        seed[(x, 0)] = Tile(N="b0", name=f"seed b0 col{x}")
    seed[(n, 0)] = Tile(N="edge", W="c1", name="seed edge")
    return seed


def _test_counter(n=6, rows=None):
    rows = rows or (1 << n) - 1
    strength = {"edge": 2}                       # value/carry glues default 1
    A, det = grow(counter_tileset(), counter_seed(n), 2, strength,
                  (0, 0, n, rows))

    def rowval(y):
        bits = []
        for x in range(n):
            t = A.get((x, y))
            if t is None:
                return None
            bits.append(1 if t.N == "b1" else 0)
        return sum(bit << (n - 1 - x) for x, bit in enumerate(bits))

    bad = filled = 0
    for y in range(1, rows + 1):
        v = rowval(y)
        if v is not None:
            filled += 1
            if v != (y & ((1 << n) - 1)):
                bad += 1
    print(f"  binary counter {n}-bit: directed={det}  rows grown={filled}  "
          f"row y encodes the integer y {'OK' if bad == 0 else f'FAIL({bad})'}")
    return det and bad == 0 and filled == rows


TILE_OUT = os.path.join(ROOT, "variants", "neural_tile.safetensors")

NBITS = 8

TAU = 2


def main_build_tile() -> int:
    ts = counter_tileset()
    strength = {"edge": 2}
    glues = sorted({g for t in ts for g in (t.N, t.E, t.S, t.W) if g})
    gid = {g: i for i, g in enumerate(glues)}
    tile_glues = torch.tensor([[gid.get(t.N, -1), gid.get(t.E, -1),
                                gid.get(t.S, -1), gid.get(t.W, -1)] for t in ts],
                              dtype=torch.long)
    glue_strength = torch.tensor([strength.get(g, 1) for g in glues], dtype=torch.long)
    # per-tile binding-gate weights = strengths of the tile's four glues (0 = null)
    bind_w = torch.tensor([[strength.get(g, 1) if g else 0 for g in (t.N, t.E, t.S, t.W)]
                           for t in ts], dtype=torch.long)
    tensors = {"tile_glues": tile_glues, "glue_strength": glue_strength,
               "binding_weight": bind_w, "binding_bias": torch.tensor(-TAU)}
    meta = {"machine": "tile", "tau": str(TAU), "glues": json.dumps(glues),
            "tile_names": json.dumps([t.name for t in ts]), "program": "binary counter"}
    save_file(tensors, TILE_OUT, metadata=meta)
    print(f"Built {os.path.relpath(TILE_OUT, ROOT)}: binary-counter tile set")
    print(f"  tiles={len(ts)}  glues={len(glues)}  tau={TAU}  size={os.path.getsize(TILE_OUT)} bytes")

    # round-trip: reconstruct the tiles from the file and regrow the counter
    t = load_file(TILE_OUT)
    with safe_open(TILE_OUT, framework="pt") as f:
        m = f.metadata()
    gl = json.loads(m["glues"])
    strg = {gl[i]: int(s) for i, s in enumerate(t["glue_strength"].tolist())}
    tiles = []
    for row, name in zip(t["tile_glues"].tolist(), json.loads(m["tile_names"])):
        sides = [gl[i] if i >= 0 else "" for i in row]
        tiles.append(Tile(N=sides[0], E=sides[1], S=sides[2], W=sides[3], name=name))

    rows = (1 << NBITS) - 1
    A, det = grow(tiles, counter_seed(NBITS), int(m["tau"]), strg,
                    (0, 0, NBITS, rows))
    bad = filled = 0
    for y in range(1, rows + 1):
        cells = [A.get((x, y)) for x in range(NBITS)]
        if any(c is None for c in cells):
            continue
        filled += 1
        v = sum((1 if c.N == "b1" else 0) << (NBITS - 1 - x) for x, c in enumerate(cells))
        if v != (y & ((1 << NBITS) - 1)):
            bad += 1
    print(f"  round-trip regrow {NBITS}-bit counter: directed={det}  rows={filled}  "
          f"row y == y {'OK' if bad == 0 else f'FAIL({bad})'}")

    # the stored binding gate reproduces the model's binding decision
    gate_ok = True
    Atest = {(1, 0): Tile(N="b0"), (2, 0): Tile(N="b0")}
    for ti, tt in enumerate(tiles):
        for site in [(1, 1), (2, 1)]:
            w = t["binding_weight"][ti].tolist()
            match = [1 if tt.glue(side) and Atest.get((site[0] + dx, site[1] + dy))
                     and tt.glue(side) == Atest[(site[0] + dx, site[1] + dy)].glue(opp) else 0
                     for dx, dy, side, opp in _SIDES]
            gate = 1 if sum(wi * mi for wi, mi in zip(w, match)) + int(t["binding_bias"]) >= 0 else 0
            if gate != int(binds(Atest, site[0], site[1], tt, TAU, strg)):
                gate_ok = False
    print(f"  stored binding gate H(weight.match - tau) matches growth rule: "
          f"{'OK' if gate_ok else 'FAIL'}")
    return 0 if (bad == 0 and det and gate_ok) else 1


def cmd_attractor():
    print("Attractor computer\n" + "=" * 40)
    print("Forward evaluation (exact, energy 0):")
    a = test_forward()
    print("Canonical relaxation:")
    b = test_energy_relax()
    print("Backward inversion (factoring by relaxation):")
    c_ = test_factor()
    print("SAT (clamp output to 1, relax to a model):")
    d = test_sat()
    print("=" * 40)
    print("ALL PASS" if (a and b and c_ and d) else "FAILURES")
    sys.exit(0 if (a and b and c_ and d) else 1)


def cmd_reversible():
    print("Reversible primitives + ALU")
    a = _test_primitives()
    b = _test_adder(4) and _test_adder(5)
    c = _test_alu(4)
    d = _test_bennett(4)
    print("PASS" if (a and b and c and d) else "FAIL")


def cmd_reversible_cpu():
    print("Reversible CPU")
    a = test_straight_line()
    b = test_bijection()
    print("PASS" if (a and b) else "FAIL")


def cmd_reversible_prog():
    print("Reversible structured programs")
    print("PASS" if _test() else "FAIL")


def cmd_ca():
    print("Reversible Margolus cellular automaton")
    a = _test_rule()
    g = _test_gates()
    b = _test_reversibility()
    c = _test_ballistic()
    d = _test_collision()
    e = _test_gate()
    f = _test_compose()
    print("PASS" if (a and g and b and c and d and e and f) else "FAIL")


def cmd_tile():
    print("Self-assembling tile computer")
    a = _test_binding_gate()
    b = all(_test_rule2(fn, nm) for fn, nm in
            [(lambda w, s: w ^ s, "XOR"), (lambda w, s: w & s, "AND"),
             (lambda w, s: w | s, "OR")])
    c = _test_counter()
    print("PASS" if (a and b and c) else "FAIL")


COMMANDS = {'attractor': cmd_attractor, 'build-attractor': main_build_attractor, 'reversible': cmd_reversible, 'reversible-cpu': cmd_reversible_cpu, 'reversible-prog': cmd_reversible_prog, 'build-reversible': main_build_reversible, 'reversible-matrix': main_reversible_matrix, 'ca': cmd_ca, 'build-ca': main_build_ca, 'tile': cmd_tile, 'build-tile': main_build_tile}


if __name__ == "__main__":
    if len(sys.argv) < 2 or sys.argv[1] not in COMMANDS:
        sys.exit('commands: ' + ', '.join(COMMANDS))
    cmd = sys.argv.pop(1)
    sys.argv[0] += ' ' + cmd
    sys.exit(COMMANDS[cmd]())