src/eval_all.py

"""
Unified evaluation harness for any threshold-computer variant.

Drops the `--cpu-test` smoke test (which was hardcoded to 16-bit/64KB) and
adds variant-aware sweep modes. The same harness handles every (data_bits,
addr_bits) configuration: it reads the manifest from each safetensors file,
runs the BatchedFitnessEvaluator at the right device, and reports per-file
plus per-category results.

Usage:
    python eval_all.py path/to/file.safetensors          # one file
    python eval_all.py variants/                          # every .safetensors in dir
    python eval_all.py --device cpu variants/             # CPU only (default)
    python eval_all.py --pop_size 32 variants/            # batched pop eval
    python eval_all.py --debug path/to/file.safetensors   # per-circuit detail
    python eval_all.py --cpu-program PATH                 # also run an assembled program
                                                          # through the threshold CPU
                                                          # sized to the file's manifest

Exit code:
    0 if all files PASS (fitness >= 0.9999)
    N where N is the number of FAILing files
"""

from __future__ import annotations

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

import torch
from safetensors import safe_open

# Reuse eval.py's evaluator (variant-aware)
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from eval import (
    BatchedFitnessEvaluator,
    create_population,
    load_model,
    get_manifest,
    heaviside,
    int_to_bits,
    bits_to_int,
    bits_msb_to_lsb,
)


# ---------------------------------------------------------------------------
# Variant-aware threshold ALU + CPU
# ---------------------------------------------------------------------------

class GenericThresholdALU:
    """Variant-aware threshold ALU walking the gates of a loaded variant.

    8-bit primitives: add8/sub8/mul8/div8, packed bitwise and8/or8/xor8,
    shl8/shr8, and the bit-cascade comparators via cmp8. Width-generic
    add_n/sub_n/mul_n/cmp_n cover the 16- and 32-bit circuit families.
    """

    def __init__(self, tensors: Dict[str, torch.Tensor], data_bits: int):
        self.T = tensors
        self.data_bits = data_bits

    def _g(self, name, inputs):
        w = self.T[name + ".weight"].view(-1)
        b = self.T[name + ".bias"].view(-1)
        return int(heaviside((torch.tensor(inputs, dtype=torch.float32) * w).sum() + b).item())

    def _xor_or_nand(self, prefix, inputs):
        a, b_ = inputs
        h_or = self._g(f"{prefix}.layer1.or", [a, b_])
        h_nand = self._g(f"{prefix}.layer1.nand", [a, b_])
        return self._g(f"{prefix}.layer2", [h_or, h_nand])

    def _fa(self, prefix, a, b, cin):
        s1 = self._xor_or_nand(f"{prefix}.ha1.sum", [a, b])
        c1 = self._g(f"{prefix}.ha1.carry", [a, b])
        s2 = self._xor_or_nand(f"{prefix}.ha2.sum", [s1, cin])
        c2 = self._g(f"{prefix}.ha2.carry", [s1, cin])
        cout = self._g(f"{prefix}.carry_or", [c1, c2])
        return s2, cout

    def add8(self, a, b):
        a_lsb = list(reversed(int_to_bits(a, 8)))
        b_lsb = list(reversed(int_to_bits(b, 8)))
        carry = 0
        s_lsb = []
        for i in range(8):
            s, carry = self._fa(f"arithmetic.ripplecarry8bit.fa{i}", a_lsb[i], b_lsb[i], carry)
            s_lsb.append(s)
        return bits_to_int(list(reversed(s_lsb))), carry

    def sub8(self, a, b):
        a_lsb = list(reversed(int_to_bits(a, 8)))
        b_lsb = list(reversed(int_to_bits(b, 8)))
        carry = 1
        d_lsb = []
        for i in range(8):
            notb = self._g(f"arithmetic.sub8bit.notb{i}", [b_lsb[i]])
            x1 = self._xor_or_nand(f"arithmetic.sub8bit.fa{i}.xor1", [a_lsb[i], notb])
            x2 = self._xor_or_nand(f"arithmetic.sub8bit.fa{i}.xor2", [x1, carry])
            and1 = self._g(f"arithmetic.sub8bit.fa{i}.and1", [a_lsb[i], notb])
            and2 = self._g(f"arithmetic.sub8bit.fa{i}.and2", [x1, carry])
            carry = self._g(f"arithmetic.sub8bit.fa{i}.or_carry", [and1, and2])
            d_lsb.append(x2)
        return bits_to_int(list(reversed(d_lsb))), carry

    _CMP_KIND = {"greaterthan": "gt", "lessthan": "lt", "eq": "eq", "equality": "eq",
                 "greaterorequal": "ge", "lessorequal": "le"}

    def _cmp_bit_cascade(self, cmp_prefix: str, finals: Dict[str, str],
                         a_bits_msb: List[int], b_bits_msb: List[int], kind: str) -> int:
        """Walk an add_bit_cascade_compare structure (bit 0 is the MSB)."""
        bits = len(a_bits_msb)
        bit_gt, bit_lt, bit_eq = [], [], []
        for i in range(bits):
            ab = [a_bits_msb[i], b_bits_msb[i]]
            bit_gt.append(self._g(f"{cmp_prefix}.bit{i}.gt", ab))
            bit_lt.append(self._g(f"{cmp_prefix}.bit{i}.lt", ab))
            eq_and = self._g(f"{cmp_prefix}.bit{i}.eq.layer1.and", ab)
            eq_nor = self._g(f"{cmp_prefix}.bit{i}.eq.layer1.nor", ab)
            bit_eq.append(self._g(f"{cmp_prefix}.bit{i}.eq", [eq_and, eq_nor]))
        cas_gt, cas_lt = [bit_gt[0]], [bit_lt[0]]
        for i in range(1, bits):
            eq_pref = self._g(f"{cmp_prefix}.cascade.eq_prefix.bit{i}", bit_eq[:i])
            cas_gt.append(self._g(f"{cmp_prefix}.cascade.gt.bit{i}", [eq_pref, bit_gt[i]]))
            cas_lt.append(self._g(f"{cmp_prefix}.cascade.lt.bit{i}", [eq_pref, bit_lt[i]]))
        if kind == "gt":
            return self._g(finals["gt"], cas_gt)
        if kind == "lt":
            return self._g(finals["lt"], cas_lt)
        if kind == "eq":
            return self._g(finals["eq"], bit_eq)
        if kind == "ge":
            lt = self._g(finals["lt"], cas_lt)
            not_lt = self._g(finals["ge"] + ".not_lt", [lt])
            return self._g(finals["ge"], [not_lt])
        if kind == "le":
            gt = self._g(finals["gt"], cas_gt)
            not_gt = self._g(finals["le"] + ".not_gt", [gt])
            return self._g(finals["le"], [not_gt])
        raise ValueError(kind)

    def cmp8(self, a, b, kind):
        return self.cmp_n(a, b, kind, 8)

    def mul8(self, a, b):
        ab = int_to_bits(a, 8)
        bb = int_to_bits(b, 8)
        result = 0
        for j in range(8):
            if bb[j] == 0:
                continue
            row = 0
            for i in range(8):
                pp = self._g(f"alu.alu8bit.mul.pp.a{i}b{j}", [ab[i], bb[j]])
                row |= (pp << (7 - i))
            shift = 7 - j
            result, _ = self.add8(result & 0xFF, (row << shift) & 0xFF)
        return result & 0xFF

    def _packed_pair_op(self, name: str, a: int, b: int) -> int:
        """8 parallel 2-input gates stored packed: weight [16], bias [8]."""
        a_bits = int_to_bits(a, 8)
        b_bits = int_to_bits(b, 8)
        w = self.T[f"{name}.weight"].view(-1)
        bias = self.T[f"{name}.bias"].view(-1)
        out = []
        for bit in range(8):
            inp = torch.tensor([float(a_bits[bit]), float(b_bits[bit])])
            out.append(int(heaviside((inp * w[bit * 2:bit * 2 + 2]).sum() + bias[bit]).item()))
        return bits_to_int(out)

    def and8(self, a, b):
        return self._packed_pair_op("alu.alu8bit.and", a, b)

    def or8(self, a, b):
        return self._packed_pair_op("alu.alu8bit.or", a, b)

    def xor8(self, a, b):
        a_bits = int_to_bits(a, 8)
        b_bits = int_to_bits(b, 8)
        w_or = self.T["alu.alu8bit.xor.layer1.or.weight"].view(-1)
        b_or = self.T["alu.alu8bit.xor.layer1.or.bias"].view(-1)
        w_nand = self.T["alu.alu8bit.xor.layer1.nand.weight"].view(-1)
        b_nand = self.T["alu.alu8bit.xor.layer1.nand.bias"].view(-1)
        w2 = self.T["alu.alu8bit.xor.layer2.weight"].view(-1)
        b2 = self.T["alu.alu8bit.xor.layer2.bias"].view(-1)
        out = []
        for bit in range(8):
            inp = torch.tensor([float(a_bits[bit]), float(b_bits[bit])])
            h_or = heaviside((inp * w_or[bit * 2:bit * 2 + 2]).sum() + b_or[bit])
            h_nand = heaviside((inp * w_nand[bit * 2:bit * 2 + 2]).sum() + b_nand[bit])
            hidden = torch.stack([h_or, h_nand])
            out.append(int(heaviside((hidden * w2[bit * 2:bit * 2 + 2]).sum() + b2[bit]).item()))
        return bits_to_int(out)

    def shl8(self, a):
        a_bits = int_to_bits(a, 8)
        out = []
        for bit in range(8):
            src = float(a_bits[bit + 1]) if bit < 7 else 0.0
            out.append(self._g(f"alu.alu8bit.shl.bit{bit}", [src]))
        return bits_to_int(out)

    def shr8(self, a):
        a_bits = int_to_bits(a, 8)
        out = []
        for bit in range(8):
            src = float(a_bits[bit - 1]) if bit > 0 else 0.0
            out.append(self._g(f"alu.alu8bit.shr.bit{bit}", [src]))
        return bits_to_int(out)

    def div8(self, a, b):
        """8-bit restoring division through the per-stage bit-cascade GE
        comparators and the sub8 subtractor. Returns (quotient, remainder);
        divide-by-zero yields (0xFF, a)."""
        if b == 0:
            return 0xFF, a
        a_bits = int_to_bits(a, 8)
        div_bits = int_to_bits(b, 8)
        quotient = 0
        remainder = 0
        for stage in range(8):
            remainder = ((remainder << 1) | a_bits[stage]) & 0xFF
            prefix = f"alu.alu8bit.div.stage{stage}"
            finals = {
                "gt": f"{prefix}.cmp_bc.gt",
                "lt": f"{prefix}.cmp_bc.lt",
                "eq": f"{prefix}.cmp_bc.eq",
                "ge": f"{prefix}.cmp",
                "le": f"{prefix}.cmp_bc.le",
            }
            ge = self._cmp_bit_cascade(f"{prefix}.cmp_bc", finals,
                                       int_to_bits(remainder, 8), div_bits, "ge")
            if ge:
                remainder, _ = self.sub8(remainder, b)
                quotient = (quotient << 1) | 1
            else:
                quotient = quotient << 1
        return quotient & 0xFF, remainder & 0xFF

    # ----- N-bit primitives (for 16-bit and 32-bit variants) ----------------

    def add_n(self, a: int, b: int, bits: int):
        """Width-generic ripple-carry add via arithmetic.ripplecarry{N}bit."""
        prefix = f"arithmetic.ripplecarry{bits}bit"
        a_lsb = list(reversed(int_to_bits(a, bits)))
        b_lsb = list(reversed(int_to_bits(b, bits)))
        carry = 0
        s_lsb = []
        for i in range(bits):
            s, carry = self._fa(f"{prefix}.fa{i}", a_lsb[i], b_lsb[i], carry)
            s_lsb.append(s)
        return bits_to_int(list(reversed(s_lsb))), carry

    def sub_n(self, a: int, b: int, bits: int):
        """N-bit two's-complement subtract via arithmetic.sub{N}bit (N >= 16).

        Structure (per build.add_sub_nbits): N NOT gates + N standard full adders.
        """
        prefix = f"arithmetic.sub{bits}bit"
        a_lsb = list(reversed(int_to_bits(a, bits)))
        b_lsb = list(reversed(int_to_bits(b, bits)))
        # NOT each B bit
        notb = [self._g(f"{prefix}.not_b.bit{i}", [b_lsb[i]]) for i in range(bits)]
        carry = 1  # carry-in = 1 for two's-complement
        d_lsb = []
        for i in range(bits):
            s, carry = self._fa(f"{prefix}.fa{i}", a_lsb[i], notb[i], carry)
            d_lsb.append(s)
        return bits_to_int(list(reversed(d_lsb))), carry

    def cmp_n(self, a: int, b: int, kind: str, bits: int):
        """N-bit unsigned comparison via the bit-cascade comparator family
        (arithmetic.cmp{N}bit.* per build.add_bit_cascade_compare)."""
        short = self._CMP_KIND[kind]
        finals = {
            "gt": f"arithmetic.greaterthan{bits}bit",
            "lt": f"arithmetic.lessthan{bits}bit",
            "eq": f"arithmetic.equality{bits}bit",
            "ge": f"arithmetic.greaterorequal{bits}bit",
            "le": f"arithmetic.lessorequal{bits}bit",
        }
        return self._cmp_bit_cascade(f"arithmetic.cmp{bits}bit", finals,
                                     int_to_bits(a, bits), int_to_bits(b, bits), short)

    def mul_n(self, a: int, b: int, bits: int):
        """N-bit shift-add multiply (low N bits only)."""
        ab = int_to_bits(a, bits)
        bb = int_to_bits(b, bits)
        mask = (1 << bits) - 1
        result = 0
        for j in range(bits):
            if bb[j] == 0:
                continue
            row = 0
            for i in range(bits):
                pp = self._g(f"alu.alu{bits}bit.mul.pp.a{i}b{j}", [ab[i], bb[j]])
                row |= (pp << (bits - 1 - i))
            shift = (bits - 1) - j
            result, _ = self.add_n(result & mask, (row << shift) & mask, bits)
        return result & mask


class GenericThresholdCPU:
    """Variant-aware CPU runtime. Sized from the variant's manifest."""

    def __init__(self, tensors: Dict[str, torch.Tensor]):
        self.T = tensors
        m = get_manifest(tensors)
        self.data_bits = m["data_bits"]
        self.addr_bits = m["addr_bits"]
        self.mem_bytes = m["memory_bytes"]
        # 8-bit CPU primitives (ripplecarry8bit, sub8bit, alu.alu8bit.*, memory.*,
        # control.*) are present in every variant regardless of manifest data_bits.
        # Wider data widths simply add additional standalone ALU primitives.
        if self.mem_bytes == 0:
            raise NotImplementedError(
                "Pure-ALU variants have no memory; cannot run CPU programs"
            )
        self.alu = GenericThresholdALU(tensors, 8)

    def _addr_decode(self, addr):
        bits = torch.tensor(int_to_bits(addr, self.addr_bits), dtype=torch.float32)
        w = self.T["memory.addr_decode.weight"]
        b = self.T["memory.addr_decode.bias"]
        return heaviside((w * bits).sum(dim=1) + b)

    def mem_read(self, mem, addr):
        sel = self._addr_decode(addr)
        mem_bits = torch.tensor(
            [int_to_bits(byte, 8) for byte in mem], dtype=torch.float32
        )
        and_w = self.T["memory.read.and.weight"]
        and_b = self.T["memory.read.and.bias"]
        or_w = self.T["memory.read.or.weight"]
        or_b = self.T["memory.read.or.bias"]
        out = []
        for bit in range(8):
            inp = torch.stack([mem_bits[:, bit], sel], dim=1)
            and_out = heaviside((inp * and_w[bit]).sum(dim=1) + and_b[bit])
            out.append(int(heaviside((and_out * or_w[bit]).sum() + or_b[bit]).item()))
        return bits_to_int(out)

    def mem_write(self, mem, addr, value):
        sel = self._addr_decode(addr)
        data_bits = torch.tensor(int_to_bits(value, 8), dtype=torch.float32)
        mem_bits = torch.tensor(
            [int_to_bits(byte, 8) for byte in mem], dtype=torch.float32
        )
        sel_w = self.T["memory.write.sel.weight"]
        sel_b = self.T["memory.write.sel.bias"]
        nsel_w = self.T["memory.write.nsel.weight"].squeeze(1)
        nsel_b = self.T["memory.write.nsel.bias"]
        and_old_w = self.T["memory.write.and_old.weight"]
        and_old_b = self.T["memory.write.and_old.bias"]
        and_new_w = self.T["memory.write.and_new.weight"]
        and_new_b = self.T["memory.write.and_new.bias"]
        or_w = self.T["memory.write.or.weight"]
        or_b = self.T["memory.write.or.bias"]
        we = torch.ones_like(sel)
        sel_inp = torch.stack([sel, we], dim=1)
        write_sel = heaviside((sel_inp * sel_w).sum(dim=1) + sel_b)
        nsel = heaviside(write_sel * nsel_w + nsel_b)
        for bit in range(8):
            old = mem_bits[:, bit]
            data_bit = data_bits[bit].expand(self.mem_bytes)
            inp_old = torch.stack([old, nsel], dim=1)
            inp_new = torch.stack([data_bit, write_sel], dim=1)
            and_old = heaviside((inp_old * and_old_w[:, bit]).sum(dim=1) + and_old_b[:, bit])
            and_new = heaviside((inp_new * and_new_w[:, bit]).sum(dim=1) + and_new_b[:, bit])
            or_inp = torch.stack([and_old, and_new], dim=1)
            new_bit = heaviside((or_inp * or_w[:, bit]).sum(dim=1) + or_b[:, bit])
            mem_bits[:, bit] = new_bit
        return [bits_to_int([int(b) for b in mem_bits[i].tolist()]) for i in range(self.mem_bytes)]

    # Conditions 0..7 map to (mux circuit, flag index into [Z, N, C, V]).
    # Odd conditions are the negated forms; the mux select is the complemented
    # flag (the circuits are plain per-bit muxes -- see control.* .inputs,
    # which name the select $not_zero / $not_carry / ... for the odd forms).
    _JCC = [("control.jz", 0), ("control.jnz", 0), ("control.jc", 2), ("control.jnc", 2),
            ("control.jn", 1), ("control.jp", 1), ("control.jv", 3), ("control.jnv", 3)]

    def _jcc_pc(self, circuit: str, pc: int, target: int, sel: int) -> int:
        """Per-bit 2:1 mux over the PC: sel ? target : pc, addr_bits wide."""
        pc_bits = int_to_bits(pc, self.addr_bits)
        t_bits = int_to_bits(target, self.addr_bits)
        out = []
        for bit in range(self.addr_bits):
            bp = f"{circuit}.bit{bit}"
            not_sel = self.alu._g(f"{bp}.not_sel", [float(sel)])
            and_a = self.alu._g(f"{bp}.and_a", [float(pc_bits[bit]), not_sel])
            and_b = self.alu._g(f"{bp}.and_b", [float(t_bits[bit]), float(sel)])
            out.append(self.alu._g(f"{bp}.or", [and_a, and_b]))
        return bits_to_int(out)

    def _sp_dec(self, sp: int) -> int:
        """SP - 1 through the control.push.sp_dec borrow chain. Gate bit index
        is MSB-first (bit addr_bits-1 is the LSB); the bit complement feeding
        each borrow AND is fixed wiring, as in the gate fitness suite."""
        bits = int_to_bits(sp, self.addr_bits)
        out = [0] * self.addr_bits
        borrow = 1
        for bit in range(self.addr_bits - 1, -1, -1):
            bp = f"control.push.sp_dec.bit{bit}"
            h_or = self.alu._g(f"{bp}.xor.layer1.or", [float(bits[bit]), float(borrow)])
            h_nand = self.alu._g(f"{bp}.xor.layer1.nand", [float(bits[bit]), float(borrow)])
            out[bit] = self.alu._g(f"{bp}.xor.layer2", [h_or, h_nand])
            borrow = self.alu._g(f"{bp}.borrow", [float(1 - bits[bit]), float(borrow)])
        return bits_to_int(out)

    def _decode_op(self, opcode):
        """4-to-16 opcode one-hot through the gate decoder."""
        ob = [float((opcode >> j) & 1) for j in range(4)]
        return [self.alu._g(f"control.decode.op{n}", ob) for n in range(16)]

    def _inc(self, pc, name, start):
        """PC + 2^start through the gate increment chain (LSB-first)."""
        ab = self.addr_bits
        pcb = [(pc >> k) & 1 for k in range(ab)]
        out = [0] * ab
        carry = 0
        for k in range(ab):
            if k < start:
                out[k] = self.alu._g(f"control.pcnext.{name}.bit{k}", [float(pcb[k])])
            else:
                cin = 1 if k == start else carry
                o = self.alu._g(f"control.pcnext.{name}.bit{k}.xor.layer1.or", [float(pcb[k]), float(cin)])
                nd = self.alu._g(f"control.pcnext.{name}.bit{k}.xor.layer1.nand", [float(pcb[k]), float(cin)])
                out[k] = self.alu._g(f"control.pcnext.{name}.bit{k}.xor.layer2", [o, nd])
                carry = self.alu._g(f"control.pcnext.{name}.bit{k}.carry", [float(pcb[k]), float(cin)])
        return sum(out[k] << k for k in range(ab))

    def _pcnext(self, pc2, pc4, addr, jcc, use_pc4, use_addr, use_jcc):
        """Next-PC priority mux through the gates: use_addr ? addr :
        (use_jcc ? jcc : (use_pc4 ? pc4 : pc2))."""
        ab = self.addr_bits
        p2 = [(pc2 >> k) & 1 for k in range(ab)]
        p4 = [(pc4 >> k) & 1 for k in range(ab)]
        aB = [(addr >> k) & 1 for k in range(ab)]
        jB = [(jcc >> k) & 1 for k in range(ab)]
        out = [0] * ab
        for k in range(ab):
            ns = self.alu._g(f"control.pcnext.m_pc4.bit{k}.not_sel", [float(use_pc4)])
            aa = self.alu._g(f"control.pcnext.m_pc4.bit{k}.and_a", [float(p2[k]), ns])
            bb = self.alu._g(f"control.pcnext.m_pc4.bit{k}.and_b", [float(p4[k]), float(use_pc4)])
            m1 = self.alu._g(f"control.pcnext.m_pc4.bit{k}.or", [aa, bb])
            ns = self.alu._g(f"control.pcnext.m_jcc.bit{k}.not_sel", [float(use_jcc)])
            aa = self.alu._g(f"control.pcnext.m_jcc.bit{k}.and_a", [float(m1), ns])
            bb = self.alu._g(f"control.pcnext.m_jcc.bit{k}.and_b", [float(jB[k]), float(use_jcc)])
            m2 = self.alu._g(f"control.pcnext.m_jcc.bit{k}.or", [aa, bb])
            ns = self.alu._g(f"control.pcnext.m_addr.bit{k}.not_sel", [float(use_addr)])
            aa = self.alu._g(f"control.pcnext.m_addr.bit{k}.and_a", [float(m2), ns])
            bb = self.alu._g(f"control.pcnext.m_addr.bit{k}.and_b", [float(aB[k]), float(use_addr)])
            out[k] = self.alu._g(f"control.pcnext.m_addr.bit{k}.or", [aa, bb])
        return sum(out[k] << k for k in range(ab))

    def step(self, state):
        if state["halted"]:
            return state
        s = dict(state)
        s["mem"] = state["mem"][:]
        s["regs"] = state["regs"][:]
        s["flags"] = state["flags"][:]
        addr_mask = (1 << self.addr_bits) - 1
        pc = s["pc"]
        hi = self.mem_read(s["mem"], pc & addr_mask)
        lo = self.mem_read(s["mem"], (pc + 1) & addr_mask)
        ir = ((hi & 0xFF) << 8) | (lo & 0xFF)
        opcode = (ir >> 12) & 0xF
        oh = self._decode_op(opcode)                 # gate opcode decode (one-hots)
        rd = (ir >> 10) & 0x3
        rs = (ir >> 8) & 0x3
        imm = ir & 0xFF
        addr = None
        if oh[0xA] or oh[0xB] or oh[0xC] or oh[0xD] or oh[0xE]:
            ah = self.mem_read(s["mem"], (pc + 2) & addr_mask)
            al = self.mem_read(s["mem"], (pc + 3) & addr_mask)
            addr = (((ah & 0xFF) << 8) | (al & 0xFF)) & addr_mask
        pc2 = self._inc(pc, "pc2", 1)                # PC+2 through the gate chain
        pc4 = self._inc(pc, "pc4", 2)                # PC+4 through the gate chain
        a = s["regs"][rd]
        b = s["regs"][rs]
        result = a
        carry = 0
        overflow = 0
        write_result = True
        jcc_val = 0
        if oh[0x0]:
            result, carry = self.alu.add8(a, b)
            overflow = 1 if (((a ^ result) & (b ^ result)) & 0x80) else 0
        elif oh[0x1]:
            result, carry = self.alu.sub8(a, b)
            overflow = 1 if (((a ^ b) & (a ^ result)) & 0x80) else 0
        elif oh[0x2]:
            result = self.alu.and8(a, b)
        elif oh[0x3]:
            result = self.alu.or8(a, b)
        elif oh[0x4]:
            result = self.alu.xor8(a, b)
        elif oh[0x5]:
            result = self.alu.shl8(a)
        elif oh[0x6]:
            result = self.alu.shr8(a)
        elif oh[0x7]:
            result = self.alu.mul8(a, b)
        elif oh[0x8]:
            result, _ = self.alu.div8(a, b)
        elif oh[0x9]:                                # CMP: set flags only
            r2, carry = self.alu.sub8(a, b)
            z = 1 if r2 == 0 else 0
            n = 1 if (r2 & 0x80) else 0
            v = 1 if (((a ^ b) & (a ^ r2)) & 0x80) else 0
            s["flags"] = [z, n, carry, v]
            write_result = False
        elif oh[0xA]:                                # LOAD
            result = self.mem_read(s["mem"], addr)
        elif oh[0xB]:                                # STORE
            s["mem"] = self.mem_write(s["mem"], addr, b & 0xFF)
            write_result = False
        elif oh[0xC]:                                # JMP (PC = addr via pcnext)
            write_result = False
        elif oh[0xD]:                                # JCC (PC via jcc mux + pcnext)
            cond = imm & 0x7
            circuit, flag_idx = self._JCC[cond]
            flag = s["flags"][flag_idx]
            sel = flag if cond % 2 == 0 else 1 - flag
            jcc_val = self._jcc_pc(circuit, pc4, addr, sel)
            write_result = False
        elif oh[0xE]:                                # CALL: push PC+4, PC = addr
            ret_addr = pc4 & 0xFFFF
            sp = s.get("sp", addr_mask)
            sp = self._sp_dec(sp)
            s["mem"] = self.mem_write(s["mem"], sp, (ret_addr >> 8) & 0xFF)
            sp = self._sp_dec(sp)
            s["mem"] = self.mem_write(s["mem"], sp, ret_addr & 0xFF)
            s["sp"] = sp
            write_result = False
        elif oh[0xF]:                                # HALT
            s["halted"] = True
            return s

        if write_result:
            s["regs"][rd] = result & 0xFF
        if oh[0x0] or oh[0x1] or oh[0x7]:
            z = 1 if (result & 0xFF) == 0 else 0
            n = 1 if (result & 0x80) else 0
            s["flags"] = [z, n, carry, overflow]
        # PC sequencing through the gate mux.
        use_pc4 = 1 if (oh[0xA] or oh[0xB]) else 0
        use_addr = 1 if (oh[0xC] or oh[0xE]) else 0
        s["pc"] = self._pcnext(pc2, pc4, addr or 0, jcc_val, use_pc4, use_addr, oh[0xD])
        return s

    def run(self, state, max_cycles=200):
        s = state
        cycles = 0
        while not s["halted"] and cycles < max_cycles:
            s = self.step(s)
            cycles += 1
        return s, cycles


def _encode_instr(opcode, rd, rs, imm):
    return ((opcode & 0xF) << 12) | ((rd & 0x3) << 10) | ((rs & 0x3) << 8) | (imm & 0xFF)


def _w16(mem, addr, value):
    mem[addr] = (value >> 8) & 0xFF
    mem[addr + 1] = value & 0xFF


PROGRAM_MIN_BYTES = 0x84  # code 0x00..0x1F + data 0x80..0x83


def builtin_program(addr_bits: int) -> Tuple[List[int], int]:
    """Sum 5+4+3+2+1 via a loop. Returns (mem, expected_result_at_0x83).

    Compact layout: code at 0x00..0x1F (32 bytes), data at 0x80..0x83 (4 bytes).
    Total footprint 132 bytes -- fits within scratchpad (256 B) and larger.
    Requires addr_bits >= 8.
    """
    if (1 << addr_bits) < PROGRAM_MIN_BYTES:
        raise ValueError(f"addr_bits={addr_bits} too small for builtin program")
    mem = [0] * (1 << addr_bits)
    mem[0x80] = 5  # initial counter
    mem[0x81] = 1  # decrement
    mem[0x82] = 0  # zero (for compare and accumulator init)
    # mem[0x83] is the output
    _w16(mem, 0x0000, _encode_instr(0xA, 1, 0, 0)); _w16(mem, 0x0002, 0x0080)
    _w16(mem, 0x0004, _encode_instr(0xA, 2, 0, 0)); _w16(mem, 0x0006, 0x0081)
    _w16(mem, 0x0008, _encode_instr(0xA, 3, 0, 0)); _w16(mem, 0x000A, 0x0082)
    _w16(mem, 0x000C, _encode_instr(0xA, 0, 0, 0)); _w16(mem, 0x000E, 0x0082)
    _w16(mem, 0x0010, _encode_instr(0x0, 0, 1, 0))
    _w16(mem, 0x0012, _encode_instr(0x1, 1, 2, 0))
    _w16(mem, 0x0014, _encode_instr(0x9, 1, 3, 0))
    _w16(mem, 0x0016, _encode_instr(0xD, 0, 0, 0x01)); _w16(mem, 0x0018, 0x0010)
    _w16(mem, 0x001A, _encode_instr(0xB, 0, 0, 0)); _w16(mem, 0x001C, 0x0083)
    _w16(mem, 0x001E, _encode_instr(0xF, 0, 0, 0))
    return mem, 15


# ---------------------------------------------------------------------------
# Eval driver
# ---------------------------------------------------------------------------

def _file_fingerprint(path: Path) -> str:
    """Stable cache key for a safetensors file: sha256 of its content.

    Hashes are content-addressed so renaming a file doesn't blow the cache,
    but mtime-only would re-key on every clone of the repo. The sha256 of a
    30 MB safetensors finishes in tens of milliseconds — small compared to
    a 5,900-test fitness run.
    """
    import hashlib
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def _cache_key(path: Path, opts: Dict[str, Any]) -> str:
    """Cache key combining file content with the relevant evaluation options."""
    fp = _file_fingerprint(path)
    opt_str = json.dumps(opts, sort_keys=True)
    import hashlib
    suffix = hashlib.sha256(opt_str.encode("utf-8")).hexdigest()[:8]
    return f"{fp}_{suffix}"


def _load_cache(cache_dir: Path, key: str) -> Dict[str, Any] | None:
    p = cache_dir / f"{key}.json"
    if not p.exists():
        return None
    try:
        return json.loads(p.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError):
        return None


def _save_cache(cache_dir: Path, key: str, payload: Dict[str, Any]) -> None:
    cache_dir.mkdir(parents=True, exist_ok=True)
    p = cache_dir / f"{key}.json"
    try:
        p.write_text(json.dumps(payload, indent=2, default=str), encoding="utf-8")
    except OSError:
        pass


def list_safetensors(path: Path) -> List[Path]:
    if path.is_file():
        return [path]
    if path.is_dir():
        return sorted(p for p in path.glob("*.safetensors") if p.is_file())
    return []


# Standalone machines carry their own circuit inventory and are verified by
# their own scripts, not the gate-fitness suite. Each maps to the tool that
# checks it.
MACHINE_VERIFIER = {
    "subleq8": "machines.py",
    "rv32": "machines.py",
    "matrix8": "matrix8.py",
    "subleq8io": "constructor8.py",
    "reflect": "reflect.py",
    "attractor": "tools/test_attractor.py",
    "reversible": "src/reversible.py",
    "ca": "src/ca.py",
    "tile": "src/tile.py",
}


def evaluate_one(path: Path, device: str, pop_size: int, debug: bool, run_cpu_program: bool) -> Dict:
    out: Dict = {"path": str(path), "filename": path.name}
    # Skip standalone machines cleanly rather than error on missing standard gates.
    with safe_open(str(path), framework="pt") as f:
        meta = f.metadata() or {}
    if meta.get("machine"):
        m = meta["machine"]
        out.update(status="SKIP", machine=m,
                   note=f"standalone machine — verify with {MACHINE_VERIFIER.get(m, 'machines.py')}")
        return out
    try:
        tensors = load_model(str(path))
    except Exception as e:
        out.update(error=f"load failed: {e}", status="ERROR")
        return out

    manifest = get_manifest(tensors)
    out.update(
        size_mb=path.stat().st_size / (1024 * 1024),
        tensors=len(tensors),
        params=sum(t.numel() for t in tensors.values()),
        manifest=manifest,
    )

    # Move to device
    tensors = {k: v.to(device) for k, v in tensors.items()}

    try:
        evaluator = BatchedFitnessEvaluator(device=device, model_path=str(path), tensors=tensors)
        population = create_population(tensors, pop_size=pop_size, device=device)
        t0 = time.perf_counter()
        fitness = evaluator.evaluate(population, debug=debug)
        elapsed = time.perf_counter() - t0
        f0 = float(fitness[0].item()) if pop_size == 1 else float(fitness.mean().item())
        out.update(
            fitness=f0,
            total_tests=evaluator.total_tests,
            elapsed_s=elapsed,
            categories={k: (float(v[0]), int(v[1])) for k, v in evaluator.category_scores.items()},
            status="PASS" if f0 >= 0.9999 else "FAIL",
        )
    except Exception as e:
        out.update(error=f"eval failed: {type(e).__name__}: {e}", status="ERROR")
        return out

    # Optional: CPU program test (8-bit CPU primitives are in every variant)
    if run_cpu_program:
        if manifest["memory_bytes"] >= PROGRAM_MIN_BYTES:
            try:
                cpu_tensors = {k: v.cpu() for k, v in tensors.items()}
                cpu = GenericThresholdCPU(cpu_tensors)
                mem, expected = builtin_program(manifest["addr_bits"])
                state = {"pc": 0, "regs": [0] * 4, "flags": [0] * 4, "mem": mem, "halted": False}
                t0 = time.perf_counter()
                final, cycles = cpu.run(state, max_cycles=200)
                cpu_elapsed = time.perf_counter() - t0
                got = final["mem"][0x83]
                out["cpu_program"] = {
                    "ok": got == expected,
                    "got": got,
                    "expected": expected,
                    "cycles": cycles,
                    "elapsed_s": cpu_elapsed,
                }
                if got != expected:
                    out["status"] = "FAIL"
            except Exception as e:
                out["cpu_program"] = {"error": str(e)}
        else:
            out["cpu_program"] = {"skipped": f"mem={manifest['memory_bytes']}B < {PROGRAM_MIN_BYTES}"}

        # Wider-ALU chain test for 16/32-bit variants
        bits = manifest["data_bits"]
        if bits in (16, 32):
            try:
                alu_tensors = {k: v.cpu() for k, v in tensors.items()}
                alu = GenericThresholdALU(alu_tensors, bits)
                t0 = time.perf_counter()
                if bits == 16:
                    x, y = 1234, 5678
                    z, _ = alu.add_n(x, y, 16);          assert z == (x + y) & 0xFFFF
                    w, _ = alu.sub_n(z, x, 16);          assert w == (z - x) & 0xFFFF, (w, z - x)
                    gt = alu.cmp_n(z, x, "greaterthan", 16); assert gt == 1
                    lt = alu.cmp_n(x, z, "lessthan",   16);  assert lt == 1
                    eq = alu.cmp_n(w, y, "eq",         16);  assert eq == 1
                    p   = alu.mul_n(123, 5, 16);          assert p == (123 * 5) & 0xFFFF
                else:  # 32
                    x, y = 1_000_000, 999_000
                    z, _ = alu.sub_n(x, y, 32);          assert z == 1_000
                    s, _ = alu.add_n(z, x, 32);          assert s == 1_001_000
                    p    = alu.mul_n(z, 100, 32);        assert p == 100_000
                    gt = alu.cmp_n(x, y, "greaterthan", 32); assert gt == 1
                    lt = alu.cmp_n(y, x, "lessthan",   32);  assert lt == 1
                    eq = alu.cmp_n(p, 100_000, "equality", 32); assert eq == 1
                chain_dt = time.perf_counter() - t0
                out[f"alu_chain_{bits}"] = {"ok": True, "elapsed_s": chain_dt}
            except AssertionError as e:
                out[f"alu_chain_{bits}"] = {"ok": False, "error": f"chain mismatch: {e}"}
                out["status"] = "FAIL"
            except Exception as e:
                out[f"alu_chain_{bits}"] = {"ok": False, "error": f"{type(e).__name__}: {e}"}
                out["status"] = "FAIL"

    return out


def print_row(r: Dict, show_cpu: bool) -> None:
    if r.get("status") == "SKIP":
        m = r.get("machine", "machine")
        print(f"  {r['filename']:<48}  SKIP  ({m} — verify with {MACHINE_VERIFIER.get(m, 'machines.py')})")
        return
    if "error" in r:
        print(f"  {r['filename']:<48}  ERROR: {r['error'][:80]}")
        return
    m = r["manifest"]
    fit = f"{r['fitness']:.4f}" if r.get("fitness") is not None else "n/a"
    cpu_col = ""
    if show_cpu and "cpu_program" in r:
        cp = r["cpu_program"]
        if cp.get("ok"):
            cpu_col = f"  CPU OK ({cp['cycles']}cyc/{cp['elapsed_s']:.1f}s)"
        elif "skipped" in cp:
            cpu_col = f"  CPU SKIP"
        elif "error" in cp:
            cpu_col = f"  CPU ERR"
        else:
            cpu_col = f"  CPU FAIL ({cp.get('got')}!={cp.get('expected')})"
    chain_col = ""
    if show_cpu:
        for bits in (16, 32):
            key = f"alu_chain_{bits}"
            if key in r:
                ch = r[key]
                if ch.get("ok"):
                    chain_col = f"  ALU{bits} OK ({ch['elapsed_s']:.2f}s)"
                else:
                    chain_col = f"  ALU{bits} FAIL"
    print(
        f"  {r['filename']:<48}  d={m['data_bits']:>2}b a={m['addr_bits']:>2}b "
        f"mem={m['memory_bytes']:>6}B  size={r['size_mb']:>6.1f}MB  "
        f"params={r['params']:>10,}  fit={fit:>6}  tests={r['total_tests']:>5}  "
        f"{r['status']:>5}{cpu_col}{chain_col}"
    )


def main() -> int:
    parser = argparse.ArgumentParser(description="Variant-agnostic eval harness")
    parser.add_argument("path", help="Path to .safetensors file or directory of files")
    parser.add_argument("--device", default="cpu", help="cpu (default) or cuda")
    parser.add_argument("--pop_size", type=int, default=1)
    parser.add_argument("--debug", action="store_true", help="Per-circuit detail per file")
    parser.add_argument("--cpu-program", action="store_true",
                        help="Also run a small assembled program through the threshold CPU "
                             "(any variant with >= 132 B memory), plus a chained 16- or "
                             "32-bit ALU sequence on wider variants")
    parser.add_argument("--json", action="store_true", help="Emit JSON results to stdout instead of a table")
    parser.add_argument("--cache-dir", default=".eval_cache",
                        help="Directory for hash-keyed result cache "
                             "(default: ./.eval_cache). Set to '' to disable.")
    parser.add_argument("--no-cache", action="store_true",
                        help="Disable the result cache for this run.")
    args = parser.parse_args()

    files = list_safetensors(Path(args.path))
    if not files:
        print(f"No .safetensors files found under {args.path}", file=sys.stderr)
        return 2

    print(f"Evaluating {len(files)} file(s) on {args.device}\n")
    cache_enabled = bool(args.cache_dir) and not args.no_cache
    cache_dir = Path(args.cache_dir) if cache_enabled else None
    cache_opts = {
        "device": args.device,
        "pop_size": args.pop_size,
        "cpu_program": bool(args.cpu_program),
    }
    cache_hits = 0
    results = []
    fail_count = 0
    for f in files:
        print(f"=== {f.name}")
        cached = None
        key = None
        if cache_enabled:
            try:
                key = _cache_key(f, cache_opts)
                cached = _load_cache(cache_dir, key)
            except OSError:
                cached = None
        if cached is not None:
            r = cached
            cache_hits += 1
            print(f"   (cache hit)")
        else:
            r = evaluate_one(f, device=args.device, pop_size=args.pop_size,
                             debug=args.debug, run_cpu_program=args.cpu_program)
            # Don't cache ERROR results: a transient failure (OOM, interrupt)
            # would otherwise be replayed on every rerun of the same file.
            if cache_enabled and key is not None and r.get("status") != "ERROR":
                _save_cache(cache_dir, key, r)
        results.append(r)
        print_row(r, show_cpu=args.cpu_program)
        if r.get("status") not in ("PASS", "SKIP"):
            fail_count += 1

    if args.json:
        # Make it serialisable
        for r in results:
            r["manifest"] = {k: (int(v) if isinstance(v, float) and v.is_integer() else v)
                             for k, v in r.get("manifest", {}).items()}
        print(json.dumps(results, indent=2, default=str))
        return fail_count

    # Summary
    print()
    print("=" * 100)
    print(" SUMMARY")
    print("=" * 100)
    for r in results:
        print_row(r, show_cpu=args.cpu_program)

    print()
    if fail_count == 0:
        print(f"ALL {len(files)} variants PASS")
    else:
        print(f"{fail_count}/{len(files)} variants FAIL")
    if cache_enabled:
        print(f"(cache: {cache_hits}/{len(files)} hits, dir={cache_dir})")
    return fail_count


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