src/eval.py

"""
Unified Evaluation Suite for 8-bit Threshold Computer
======================================================
GPU-batched evaluation with per-circuit reporting.
Includes CPU runtime for threshold-weight execution.

Usage:
    python eval.py                    # Run circuit evaluation
    python eval.py --device cpu       # CPU mode
    python eval.py --pop_size 1000    # Population mode for evolution
    python eval.py --cpu-test         # Run CPU smoke test

The single gate-routed CPU runtime is GenericThresholdCPU in eval_all.py
(manifest-sized); the pure-Python reference (CPUState / ref_step) lives here
for cross-checks.

API (for prune_weights.py):
    from eval import load_model, create_population, BatchedFitnessEvaluator
"""

import argparse
import json
import os
import time
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Callable, Dict, List, Optional, Tuple

import torch
from safetensors import safe_open


MODEL_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "neural_computer.safetensors")  # repo root; this module lives in src/


@dataclass
class CircuitResult:
    """Result for a single circuit test."""
    name: str
    passed: int
    total: int
    failures: List[Tuple] = field(default_factory=list)

    @property
    def success(self) -> bool:
        return self.passed == self.total

    @property
    def rate(self) -> float:
        return self.passed / self.total if self.total > 0 else 0.0


def heaviside(x: torch.Tensor) -> torch.Tensor:
    """Threshold activation: 1 if x >= 0, else 0."""
    return (x >= 0).float()


def load_model(path: str = MODEL_PATH) -> Dict[str, torch.Tensor]:
    """Load model tensors from safetensors."""
    with safe_open(path, framework='pt') as f:
        return {name: f.get_tensor(name).float() for name in f.keys()}


def load_metadata(path: str = MODEL_PATH) -> Dict:
    """Load metadata from safetensors (includes signal_registry)."""
    with safe_open(path, framework='pt') as f:
        meta = f.metadata()
        if meta and 'signal_registry' in meta:
            return {'signal_registry': json.loads(meta['signal_registry'])}
        return {'signal_registry': {}}


def get_manifest(tensors: Dict[str, torch.Tensor]) -> Dict[str, int]:
    """Extract manifest values from tensors.

    Returns dict with data_bits, addr_bits, memory_bytes, version.
    Defaults to 8-bit data, 16-bit addr for legacy models.
    """
    return {
        'data_bits': int(tensors.get('manifest.data_bits', torch.tensor([8.0])).item()),
        'addr_bits': int(tensors.get('manifest.addr_bits',
                         tensors.get('manifest.pc_width', torch.tensor([16.0]))).item()),
        'memory_bytes': int(tensors.get('manifest.memory_bytes', torch.tensor([65536.0])).item()),
        'version': float(tensors.get('manifest.version', torch.tensor([1.0])).item()),
    }


def create_population(
    base_tensors: Dict[str, torch.Tensor],
    pop_size: int,
    device: str = 'cuda'
) -> Dict[str, torch.Tensor]:
    """Replicate base tensors for batched population evaluation."""
    return {
        name: tensor.unsqueeze(0).expand(pop_size, *tensor.shape).clone().to(device)
        for name, tensor in base_tensors.items()
    }


# =============================================================================
# CPU RUNTIME
# =============================================================================

FLAG_NAMES = ["Z", "N", "C", "V"]
CTRL_NAMES = ["HALT", "MEM_WE", "MEM_RE", "RESERVED"]

PC_BITS = 16
IR_BITS = 16
REG_BITS = 8
REG_COUNT = 4
FLAG_BITS = 4
SP_BITS = 16
CTRL_BITS = 4
MEM_BYTES = 65536
MEM_BITS = MEM_BYTES * 8

STATE_BITS = PC_BITS + IR_BITS + (REG_BITS * REG_COUNT) + FLAG_BITS + SP_BITS + CTRL_BITS + MEM_BITS


def int_to_bits(value: int, width: int) -> List[int]:
    return [(value >> (width - 1 - i)) & 1 for i in range(width)]


def bits_to_int(bits: List[int]) -> int:
    value = 0
    for bit in bits:
        value = (value << 1) | int(bit)
    return value


def bits_msb_to_lsb(bits: List[int]) -> List[int]:
    return list(reversed(bits))


@dataclass
class CPUState:
    pc: int
    ir: int
    regs: List[int]
    flags: List[int]
    sp: int
    ctrl: List[int]
    mem: List[int]

    def copy(self) -> 'CPUState':
        return CPUState(
            pc=int(self.pc),
            ir=int(self.ir),
            regs=[int(r) for r in self.regs],
            flags=[int(f) for f in self.flags],
            sp=int(self.sp),
            ctrl=[int(c) for c in self.ctrl],
            mem=[int(m) for m in self.mem],
        )


def pack_state(state: CPUState) -> List[int]:
    bits: List[int] = []
    bits.extend(int_to_bits(state.pc, PC_BITS))
    bits.extend(int_to_bits(state.ir, IR_BITS))
    for reg in state.regs:
        bits.extend(int_to_bits(reg, REG_BITS))
    bits.extend([int(f) for f in state.flags])
    bits.extend(int_to_bits(state.sp, SP_BITS))
    bits.extend([int(c) for c in state.ctrl])
    for byte in state.mem:
        bits.extend(int_to_bits(byte, REG_BITS))
    return bits


def unpack_state(bits: List[int]) -> CPUState:
    if len(bits) != STATE_BITS:
        raise ValueError(f"Expected {STATE_BITS} bits, got {len(bits)}")

    idx = 0
    pc = bits_to_int(bits[idx:idx + PC_BITS])
    idx += PC_BITS
    ir = bits_to_int(bits[idx:idx + IR_BITS])
    idx += IR_BITS

    regs = []
    for _ in range(REG_COUNT):
        regs.append(bits_to_int(bits[idx:idx + REG_BITS]))
        idx += REG_BITS

    flags = [int(b) for b in bits[idx:idx + FLAG_BITS]]
    idx += FLAG_BITS

    sp = bits_to_int(bits[idx:idx + SP_BITS])
    idx += SP_BITS

    ctrl = [int(b) for b in bits[idx:idx + CTRL_BITS]]
    idx += CTRL_BITS

    mem = []
    for _ in range(MEM_BYTES):
        mem.append(bits_to_int(bits[idx:idx + REG_BITS]))
        idx += REG_BITS

    return CPUState(pc=pc, ir=ir, regs=regs, flags=flags, sp=sp, ctrl=ctrl, mem=mem)


def decode_ir(ir: int) -> Tuple[int, int, int, int]:
    opcode = (ir >> 12) & 0xF
    rd = (ir >> 10) & 0x3
    rs = (ir >> 8) & 0x3
    imm8 = ir & 0xFF
    return opcode, rd, rs, imm8


def flags_from_result(result: int, carry: int, overflow: int) -> Tuple[int, int, int, int]:
    z = 1 if result == 0 else 0
    n = 1 if (result & 0x80) else 0
    c = 1 if carry else 0
    v = 1 if overflow else 0
    return z, n, c, v


def alu_add(a: int, b: int) -> Tuple[int, int, int]:
    full = a + b
    result = full & 0xFF
    carry = 1 if full > 0xFF else 0
    overflow = 1 if (((a ^ result) & (b ^ result)) & 0x80) else 0
    return result, carry, overflow


def alu_sub(a: int, b: int) -> Tuple[int, int, int]:
    full = (a - b) & 0x1FF
    result = full & 0xFF
    carry = 1 if a >= b else 0
    overflow = 1 if (((a ^ b) & (a ^ result)) & 0x80) else 0
    return result, carry, overflow


def ref_step(state: CPUState) -> CPUState:
    """Reference CPU cycle (pure Python arithmetic)."""
    if state.ctrl[0] == 1:
        return state.copy()

    s = state.copy()

    hi = s.mem[s.pc]
    lo = s.mem[(s.pc + 1) & 0xFFFF]
    s.ir = ((hi & 0xFF) << 8) | (lo & 0xFF)
    next_pc = (s.pc + 2) & 0xFFFF

    opcode, rd, rs, imm8 = decode_ir(s.ir)
    a = s.regs[rd]
    b = s.regs[rs]

    addr16 = None
    next_pc_ext = next_pc
    if opcode in (0xA, 0xB, 0xC, 0xD, 0xE):
        addr_hi = s.mem[next_pc]
        addr_lo = s.mem[(next_pc + 1) & 0xFFFF]
        addr16 = ((addr_hi & 0xFF) << 8) | (addr_lo & 0xFF)
        next_pc_ext = (next_pc + 2) & 0xFFFF

    write_result = True
    result = a
    carry = 0
    overflow = 0

    if opcode == 0x0:
        result, carry, overflow = alu_add(a, b)
    elif opcode == 0x1:
        result, carry, overflow = alu_sub(a, b)
    elif opcode == 0x2:
        result = a & b
    elif opcode == 0x3:
        result = a | b
    elif opcode == 0x4:
        result = a ^ b
    elif opcode == 0x5:
        result = (a << 1) & 0xFF
    elif opcode == 0x6:
        result = (a >> 1) & 0xFF
    elif opcode == 0x7:
        result = (a * b) & 0xFF
    elif opcode == 0x8:
        if b == 0:
            result = 0xFF
        else:
            result = a // b
    elif opcode == 0x9:
        result, carry, overflow = alu_sub(a, b)
        write_result = False
    elif opcode == 0xA:
        result = s.mem[addr16]
    elif opcode == 0xB:
        s.mem[addr16] = b & 0xFF
        write_result = False
    elif opcode == 0xC:
        s.pc = addr16 & 0xFFFF
        write_result = False
    elif opcode == 0xD:
        cond_type = imm8 & 0x7
        if cond_type == 0:
            take_branch = s.flags[0] == 1
        elif cond_type == 1:
            take_branch = s.flags[0] == 0
        elif cond_type == 2:
            take_branch = s.flags[2] == 1
        elif cond_type == 3:
            take_branch = s.flags[2] == 0
        elif cond_type == 4:
            take_branch = s.flags[1] == 1
        elif cond_type == 5:
            take_branch = s.flags[1] == 0
        elif cond_type == 6:
            take_branch = s.flags[3] == 1
        else:
            take_branch = s.flags[3] == 0
        if take_branch:
            s.pc = addr16 & 0xFFFF
        else:
            s.pc = next_pc_ext
        write_result = False
    elif opcode == 0xE:
        ret_addr = next_pc_ext & 0xFFFF
        s.sp = (s.sp - 1) & 0xFFFF
        s.mem[s.sp] = (ret_addr >> 8) & 0xFF
        s.sp = (s.sp - 1) & 0xFFFF
        s.mem[s.sp] = ret_addr & 0xFF
        s.pc = addr16 & 0xFFFF
        write_result = False
    elif opcode == 0xF:
        s.ctrl[0] = 1
        write_result = False

    # Flag policy: only ADD, SUB, MUL, and CMP write Z/N/C/V (MUL clears C
    # and V). Bitwise, shift, DIV, LOAD, and STORE leave FLAGS unchanged.
    if opcode in (0x0, 0x1, 0x7, 0x9):
        s.flags = list(flags_from_result(result, carry, overflow))

    if write_result:
        s.regs[rd] = result & 0xFF

    if opcode not in (0xC, 0xD, 0xE):
        s.pc = next_pc_ext

    return s


def ref_run_until_halt(state: CPUState, max_cycles: int = 256) -> Tuple[CPUState, int]:
    """Reference execution loop."""
    s = state.copy()
    for i in range(max_cycles):
        if s.ctrl[0] == 1:
            return s, i
        s = ref_step(s)
    return s, max_cycles


def encode_instr(opcode: int, rd: int, rs: int, imm8: int) -> int:
    return ((opcode & 0xF) << 12) | ((rd & 0x3) << 10) | ((rs & 0x3) << 8) | (imm8 & 0xFF)


def write_instr(mem: List[int], addr: int, instr: int) -> None:
    mem[addr & 0xFFFF] = (instr >> 8) & 0xFF
    mem[(addr + 1) & 0xFFFF] = instr & 0xFF


def write_addr(mem: List[int], addr: int, value: int) -> None:
    mem[addr & 0xFFFF] = (value >> 8) & 0xFF
    mem[(addr + 1) & 0xFFFF] = value & 0xFF


def _fill_smoke_program(mem: List[int]) -> None:
    """LOAD/ADD/STORE, MUL with a high-bit operand, and a SUB/JNZ loop."""
    write_instr(mem, 0x0000, encode_instr(0xA, 0, 0, 0x00))
    write_addr(mem, 0x0002, 0x0100)
    write_instr(mem, 0x0004, encode_instr(0xA, 1, 0, 0x00))
    write_addr(mem, 0x0006, 0x0101)
    write_instr(mem, 0x0008, encode_instr(0x0, 0, 1, 0x00))
    write_instr(mem, 0x000A, encode_instr(0xB, 0, 0, 0x00))
    write_addr(mem, 0x000C, 0x0102)
    write_instr(mem, 0x000E, encode_instr(0xA, 2, 0, 0x00))
    write_addr(mem, 0x0010, 0x0103)
    write_instr(mem, 0x0012, encode_instr(0xA, 3, 0, 0x00))
    write_addr(mem, 0x0014, 0x0104)
    write_instr(mem, 0x0016, encode_instr(0x7, 2, 3, 0x00))
    write_instr(mem, 0x0018, encode_instr(0xB, 0, 2, 0x00))
    write_addr(mem, 0x001A, 0x0105)
    write_instr(mem, 0x001C, encode_instr(0xA, 0, 0, 0x00))
    write_addr(mem, 0x001E, 0x0106)
    write_instr(mem, 0x0020, encode_instr(0xA, 1, 0, 0x00))
    write_addr(mem, 0x0022, 0x0107)
    write_instr(mem, 0x0024, encode_instr(0xA, 3, 0, 0x00))
    write_addr(mem, 0x0026, 0x0108)
    write_instr(mem, 0x0028, encode_instr(0x0, 1, 0, 0x00))   # loop: R1 += R0
    write_instr(mem, 0x002A, encode_instr(0x1, 0, 3, 0x00))   #       R0 -= 1
    write_instr(mem, 0x002C, encode_instr(0xD, 0, 0, 0x01))   #       JNZ loop
    write_addr(mem, 0x002E, 0x0028)
    write_instr(mem, 0x0030, encode_instr(0xB, 0, 1, 0x00))
    write_addr(mem, 0x0032, 0x0109)
    write_instr(mem, 0x0034, encode_instr(0xF, 0, 0, 0x00))
    mem[0x0100] = 5
    mem[0x0101] = 7
    mem[0x0103] = 2
    mem[0x0104] = 131
    mem[0x0106] = 3
    mem[0x0107] = 0
    mem[0x0108] = 1


def run_smoke_test() -> int:
    """Smoke test through the single gate-routed CPU runtime
    (GenericThresholdCPU), cross-checked against the pure-Python reference:

    1. LOAD 5, LOAD 7, ADD, STORE       -> MEM[0x0102] = 12
    2. LOAD 2, LOAD 131, MUL, STORE     -> MEM[0x0105] = (2*131) & 0xFF = 6
    3. Countdown loop 3+2+1 via SUB/JNZ -> MEM[0x0109] = 6

    Runs on the 1 KB variant so the threshold pass finishes in seconds; the
    reference runs at 64 KB in pure Python.
    """
    import os
    from eval_all import GenericThresholdCPU   # lazy: avoids eval<->eval_all cycle

    expected = {0x0102: 12, 0x0105: 6, 0x0109: 6}

    print("Running reference implementation...")
    ref_mem = [0] * 65536
    _fill_smoke_program(ref_mem)
    state = CPUState(pc=0, ir=0, regs=[0, 0, 0, 0], flags=[0, 0, 0, 0],
                     sp=0xFFFE, ctrl=[0, 0, 0, 0], mem=ref_mem)
    final, cycles = ref_run_until_halt(state, max_cycles=40)
    assert final.ctrl[0] == 1, "HALT flag not set"
    for addr, want in expected.items():
        assert final.mem[addr] == want, f"MEM[{addr:#06x}] expected {want}, got {final.mem[addr]}"
    print(f"  Reference: ADD={final.mem[0x0102]}, MUL={final.mem[0x0105]}, "
          f"LOOP={final.mem[0x0109]}, cycles={cycles}")

    print("Running threshold-weight implementation (GenericThresholdCPU, 1 KB)...")
    path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "variants",
                        "neural_computer8_small.safetensors")
    tensors = {}
    with safe_open(path, framework="pt") as f:
        for name in f.keys():
            tensors[name] = f.get_tensor(name).float()
    cpu = GenericThresholdCPU(tensors)
    t_mem = [0] * 1024
    _fill_smoke_program(t_mem)
    t_state = {"pc": 0, "regs": [0] * 4, "flags": [0] * 4, "mem": t_mem, "halted": False}
    t_final, t_cycles = cpu.run(t_state, max_cycles=40)
    assert t_final["halted"], "Threshold HALT not reached"
    for addr, want in expected.items():
        assert t_final["mem"][addr] == want, (
            f"Threshold MEM[{addr:#06x}] mismatch: {t_final['mem'][addr]} != {want}")
    assert t_cycles == cycles, f"cycle count mismatch: {t_cycles} != {cycles}"
    print(f"  Threshold: ADD={t_final['mem'][0x0102]}, MUL={t_final['mem'][0x0105]}, "
          f"LOOP={t_final['mem'][0x0109]}, cycles={t_cycles}")

    print("\nSmoke test: PASSED")
    return 0


# =============================================================================
# NETLIST EVALUATION (metadata-driven)
# =============================================================================

class NetlistEvaluator:
    """Evaluate a self-contained circuit from its shipped wiring metadata.

    The gate graph is built from the .inputs tensors and the header signal
    registry — the same artifacts safetensors2verilog consumes — so a test
    driven through this class proves the file itself encodes the circuit,
    with no wiring knowledge living in Python.

    External signals are names starting with '$' (plus the constants #0/#1);
    the caller binds them per evaluation. Gate outputs are published under
    their full gate names. Evaluation is batched over test vectors and,
    for population dicts, over population slots.
    """

    def __init__(self, tensors: Dict[str, torch.Tensor],
                 signal_registry: Dict[str, str], prefix: str,
                 pop_size: int = 1, levels: bool = True):
        id_to_name = {int(k): v for k, v in signal_registry.items()}
        self.prefix = prefix
        self.pop_size = pop_size
        self.gates: Dict[str, Tuple[torch.Tensor, torch.Tensor, List[str]]] = {}
        for key, t in tensors.items():
            if not key.endswith('.inputs'):
                continue
            gate = key[: -len('.inputs')]
            if gate != prefix and not gate.startswith(prefix + '.'):
                continue
            w = tensors.get(gate + '.weight')
            b = tensors.get(gate + '.bias')
            if w is None or b is None:
                continue
            inp = t.reshape(pop_size, -1)[0] if pop_size > 1 else t.flatten()
            fan = inp.numel()
            if w.numel() != fan * pop_size:
                continue  # packed multi-gate tensor; not a single netlist gate
            names = [id_to_name[int(i)] for i in inp.tolist()]
            self.gates[gate] = (w.float().view(pop_size, fan),
                                b.float().view(pop_size), names)
        if not self.gates:
            raise KeyError(f"no wired gates under prefix {prefix}")

        # Topological order (Kahn). Inputs that are gates in this circuit are
        # dependencies; everything else must be bound externally.
        indeg = {g: 0 for g in self.gates}
        consumers: Dict[str, List[str]] = {}
        for g, (_, _, names) in self.gates.items():
            for n in names:
                if n in self.gates:
                    indeg[g] += 1
                    consumers.setdefault(n, []).append(g)
        order = [g for g, d in indeg.items() if d == 0]
        i = 0
        while i < len(order):
            for c in consumers.get(order[i], []):
                indeg[c] -= 1
                if indeg[c] == 0:
                    order.append(c)
            i += 1
        if len(order) != len(self.gates):
            cyc = sorted(set(self.gates) - set(order))[:5]
            raise ValueError(f"wiring cycle under {prefix}: {cyc}")
        self.order = order

        # Leveled plan: one padded tensor op per topological level instead of
        # one Python step per gate. Turns thousands of tiny ops into ~depth
        # batched ops (a MULH-class netlist drops from seconds to tens of ms).
        self._plan = None
        if levels:
            self._build_levels(id_to_name)

    def _build_levels(self, id_to_name):
        # Signal slots: constants, all externals, all gate outputs.
        ext = self.external_names()
        names = ['#0', '#1'] + ext + self.order
        self.slot = {n: i for i, n in enumerate(names)}
        self.n_sig = len(names)
        self.ext_names = ext

        depth = {}
        for n in names:
            if n not in self.gates:
                depth[n] = 0
        for g in self.order:
            _, _, ins = self.gates[g]
            depth[g] = 1 + max((depth[n] for n in ins), default=0)
        by_level: Dict[int, List[str]] = {}
        for g in self.order:
            by_level.setdefault(depth[g], []).append(g)

        plan = []
        for lvl in sorted(by_level):
            gs = by_level[lvl]
            max_fan = max(len(self.gates[g][2]) for g in gs)
            n_g = len(gs)
            idx = torch.zeros(n_g, max_fan, dtype=torch.long)
            w = torch.zeros(n_g, max_fan, self.pop_size)
            b = torch.zeros(n_g, self.pop_size)
            out = torch.zeros(n_g, dtype=torch.long)
            for r, g in enumerate(gs):
                wt, bs, ins = self.gates[g]
                out[r] = self.slot[g]
                b[r] = bs
                for c, n in enumerate(ins):
                    idx[r, c] = self.slot[n]
                    w[r, c] = wt[:, c]
            plan.append((idx, w, b, out))
        self._plan = plan

    def external_names(self) -> List[str]:
        """Every non-gate, non-constant signal the circuit consumes."""
        out = set()
        for _, (_, _, names) in self.gates.items():
            for n in names:
                if n not in self.gates and n not in ('#0', '#1'):
                    out.add(n)
        return sorted(out)

    def run(self, external: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        """Evaluate the circuit. Each external is a [num_tests] tensor (or a
        scalar); every returned gate output is [num_tests, pop_size]."""
        first = next(iter(external.values()))
        if not torch.is_tensor(first):
            num_tests = 1
        else:
            num_tests = first.shape[0] if first.dim() else 1
        device = self.gates[self.order[0]][0].device

        def expand(v):
            t = v if torch.is_tensor(v) else torch.tensor(float(v))
            t = t.float().to(device).reshape(-1)
            return t.unsqueeze(1).expand(num_tests, self.pop_size)

        if self._plan is not None:
            V = torch.zeros(self.n_sig, num_tests, self.pop_size, device=device)
            V[self.slot['#1']] = 1.0
            for k, v in external.items():
                V[self.slot[k]] = expand(v)
            for idx, w, b, out in self._plan:
                # gathered: [n_g, max_fan, num_tests, pop]
                gathered = V[idx]
                acc = (gathered * w[:, :, None, :]).sum(1) + b[:, None, :]
                V[out] = (acc >= 0).float()
            return _SlotView(V, self.slot)

        values: Dict[str, torch.Tensor] = {
            '#0': torch.zeros(num_tests, self.pop_size, device=device),
            '#1': torch.ones(num_tests, self.pop_size, device=device),
        }
        for k, v in external.items():
            values[k] = expand(v)
        for g in self.order:
            w, b, names = self.gates[g]
            acc = b.unsqueeze(0).expand(num_tests, self.pop_size).clone()
            for k, n in enumerate(names):
                acc = acc + w[:, k] * values[n]
            values[g] = (acc >= 0).float()
        return values


class _SlotView:
    """Read gate outputs by name from a packed [n_sig, num_tests, pop] tensor."""
    __slots__ = ("_V", "_slot")

    def __init__(self, V, slot):
        self._V = V
        self._slot = slot

    def __getitem__(self, name):
        return self._V[self._slot[name]]

    def __contains__(self, name):
        return name in self._slot


def float_bits_to_value(word: int, exp_bits: int, frac_bits: int) -> float:
    """Decode an IEEE 754 word to a Python float. float64 represents every
    float16/float32 value exactly, so the oracle comparisons are exact."""
    sign = (word >> (exp_bits + frac_bits)) & 1
    exp = (word >> frac_bits) & ((1 << exp_bits) - 1)
    frac = word & ((1 << frac_bits) - 1)
    bias = (1 << (exp_bits - 1)) - 1
    if exp == (1 << exp_bits) - 1:
        if frac:
            return float('nan')
        return float('-inf') if sign else float('inf')
    if exp == 0:
        v = frac * 2.0 ** (1 - bias - frac_bits)
    else:
        v = (frac + (1 << frac_bits)) * 2.0 ** (exp - bias - frac_bits)
    return -v if sign else v


def float_test_words(exp_bits: int, frac_bits: int) -> Tuple[List[int], List[int]]:
    """Directed IEEE edge encodings (every category, both signs) plus seeded
    random words for a family."""
    import random
    E, F = exp_bits, frac_bits
    emax = (1 << E) - 1
    fmax = (1 << F) - 1
    bias = (1 << (E - 1)) - 1

    def word(s, e, f):
        return (s << (E + F)) | (e << F) | f

    directed = []
    for s in (0, 1):
        directed += [
            word(s, 0, 0),                # +-0
            word(s, 0, 1),                # smallest subnormal
            word(s, 0, fmax),             # largest subnormal
            word(s, 1, 0),                # smallest normal
            word(s, bias, 0),             # +-1.0
            word(s, bias, 1 << (F - 1)),  # +-1.5
            word(s, emax - 1, fmax),      # largest normal
            word(s, emax, 0),             # +-inf
            word(s, emax, 1),             # NaN, minimal payload
            word(s, emax, fmax),          # NaN, full payload
        ]
    rng = random.Random(0xF10A7 + E)
    randoms = [rng.getrandbits(1 + E + F) for _ in range(24)]
    return directed, randoms


def float_mul_oracle(aw: int, bw: int, exp_bits: int, frac_bits: int) -> int:
    """Expected product word under the documented contract: exact IEEE
    specials (NaN, infinities, signed zeros), subnormal operands and
    gradual-underflow subnormal results, round-to-nearest-even mantissa.
    Pure integer arithmetic, so exact."""
    E, F = exp_bits, frac_bits
    emax = (1 << E) - 1
    fmask = (1 << F) - 1
    bias = (1 << (E - 1)) - 1
    qnan = (emax << F) | (1 << (F - 1))
    sa, ea, fa = (aw >> (E + F)) & 1, (aw >> F) & emax, aw & fmask
    sb, eb, fb = (bw >> (E + F)) & 1, (bw >> F) & emax, bw & fmask
    s = sa ^ sb
    a_nan = ea == emax and fa != 0
    b_nan = eb == emax and fb != 0
    a_inf = ea == emax and fa == 0
    b_inf = eb == emax and fb == 0
    a_zero = ea == 0 and fa == 0
    b_zero = eb == 0 and fb == 0
    if a_nan or b_nan or (a_inf and b_zero) or (b_inf and a_zero):
        return qnan
    if a_inf or b_inf:
        return (s << (E + F)) | (emax << F)
    if a_zero or b_zero:
        return s << (E + F)
    Ma = ((1 << F) if ea else 0) | fa
    Mb = ((1 << F) if eb else 0) | fb
    eea = ea if ea else 1
    eeb = eb if eb else 1
    P = Ma * Mb
    t = P.bit_length() - 1
    er = t + eea + eeb - bias - 2 * F
    if er >= 1:
        drop = t - F
        if drop <= 0:
            mant = P << (-drop)
            guard = sticky = 0
        else:
            mant = P >> drop
            guard = (P >> (drop - 1)) & 1
            sticky = 1 if (P & ((1 << (drop - 1)) - 1)) else 0
        frac = mant & fmask
        if guard and ((frac & 1) or sticky):
            frac += 1
            if frac > fmask:
                frac = 0
                er += 1
        if er >= emax:
            return (s << (E + F)) | (emax << F)
        return (s << (E + F)) | (er << F) | frac
    sh_q = eea + eeb - bias - F - 1               # subnormal: value in quanta
    if sh_q >= 0:
        sig = P << sh_q
        guard = sticky = 0
    else:
        r = -sh_q
        sig = P >> r
        guard = (P >> (r - 1)) & 1
        sticky = 1 if (P & ((1 << (r - 1)) - 1)) else 0
    if guard and ((sig & 1) or sticky):
        sig += 1
    if sig == 0:
        return s << (E + F)
    if sig >= (1 << F):
        return (s << (E + F)) | (1 << F)
    return (s << (E + F)) | sig


def float_div_oracle(aw: int, bw: int, exp_bits: int, frac_bits: int) -> int:
    """Expected quotient word under the documented contract: exact IEEE
    specials (NaN, infinities, x/0 -> inf, 0/0 and inf/inf -> NaN, signed
    zeros), subnormal operands and gradual-underflow subnormal results,
    round-to-nearest-even mantissa."""
    E, F = exp_bits, frac_bits
    emax = (1 << E) - 1
    fmask = (1 << F) - 1
    bias = (1 << (E - 1)) - 1
    qnan = (emax << F) | (1 << (F - 1))
    sa, ea, fa = (aw >> (E + F)) & 1, (aw >> F) & emax, aw & fmask
    sb, eb, fb = (bw >> (E + F)) & 1, (bw >> F) & emax, bw & fmask
    s = sa ^ sb
    a_nan = ea == emax and fa != 0
    b_nan = eb == emax and fb != 0
    a_inf = ea == emax and fa == 0
    b_inf = eb == emax and fb == 0
    a_zero = ea == 0 and fa == 0
    b_zero = eb == 0 and fb == 0
    if a_nan or b_nan or (a_zero and b_zero) or (a_inf and b_inf):
        return qnan
    if a_inf or b_zero:
        return (s << (E + F)) | (emax << F)
    if a_zero or b_inf:
        return s << (E + F)
    Ma = ((1 << F) if ea else 0) | fa
    Mb = ((1 << F) if eb else 0) | fb
    eea = ea if ea else 1
    eeb = eb if eb else 1
    ta = Ma.bit_length() - 1                       # pre-normalize both mantissas
    tb = Mb.bit_length() - 1
    na = eea - (F - ta)
    nb = eeb - (F - tb)
    Ma <<= (F - ta)
    Mb <<= (F - tb)
    q0 = 1 if Ma >= Mb else 0
    er = na - nb + bias - 1 + q0
    num = Ma << (F + 2 - q0)
    Q = num // Mb
    rem = num % Mb
    if er >= 1:
        frac = (Q >> 1) & fmask
        guard = Q & 1
        sticky = 1 if rem else 0
        if guard and ((frac & 1) or sticky):
            frac += 1
            if frac > fmask:
                frac = 0
                er += 1
        if er >= emax:
            return (s << (E + F)) | (emax << F)
        return (s << (E + F)) | (er << F) | frac
    rsh = (1 - er) + 1                              # gradual underflow
    sig = Q >> rsh
    guard = (Q >> (rsh - 1)) & 1
    sticky = 1 if ((Q & ((1 << (rsh - 1)) - 1)) or rem) else 0
    if guard and ((sig & 1) or sticky):
        sig += 1
    if sig == 0:
        return s << (E + F)
    if sig >= (1 << F):
        return (s << (E + F)) | (1 << F)
    return (s << (E + F)) | sig


def float_fma_oracle(aw: int, bw: int, cw: int, exp_bits: int, frac_bits: int) -> int:
    """Expected word for round(a*b + c) with a SINGLE rounding: exact IEEE
    specials (NaN, infinities, inf*0 and inf-inf -> NaN, signed zeros, exact
    cancellation -> +0), subnormal operands and gradual-underflow results,
    round-to-nearest-even. Pure integer arithmetic, so exact."""
    E, F = exp_bits, frac_bits
    emax = (1 << E) - 1
    fmask = (1 << F) - 1
    bias = (1 << (E - 1)) - 1
    qnan = (emax << F) | (1 << (F - 1))

    def dec(w):
        return (w >> (E + F)) & 1, (w >> F) & emax, w & fmask
    sa, ea, fa = dec(aw)
    sb, eb, fb = dec(bw)
    sc, ec, fc = dec(cw)
    sp = sa ^ sb
    a_nan = ea == emax and fa != 0
    b_nan = eb == emax and fb != 0
    c_nan = ec == emax and fc != 0
    a_inf = ea == emax and fa == 0
    b_inf = eb == emax and fb == 0
    c_inf = ec == emax and fc == 0
    a_zero = ea == 0 and fa == 0
    b_zero = eb == 0 and fb == 0
    c_zero = ec == 0 and fc == 0
    if a_nan or b_nan or c_nan:
        return qnan
    if (a_inf and b_zero) or (b_inf and a_zero):
        return qnan
    if a_inf or b_inf:
        if c_inf and sc != sp:
            return qnan
        return (sp << (E + F)) | (emax << F)
    if c_inf:
        return (sc << (E + F)) | (emax << F)

    Ma = ((1 << F) if ea else 0) | fa
    Mb = ((1 << F) if eb else 0) | fb
    Mc = ((1 << F) if ec else 0) | fc
    eea = ea if ea else 1
    eeb = eb if eb else 1
    eec = ec if ec else 1
    P = Ma * Mb
    pe = eea + eeb - 2 * bias - 2 * F
    ce = eec - bias - F
    if P == 0 and c_zero:
        return ((1 if (sp and sc) else 0) << (E + F))
    g = min(pe, ce)
    Pi = P << (pe - g)
    Ci = Mc << (ce - g)
    S = (-Pi if sp else Pi) + (-Ci if sc else Ci)
    if S == 0:
        return 0
    s_out = 1 if S < 0 else 0
    mag = -S if S < 0 else S
    t = mag.bit_length() - 1
    er = t + g + bias
    if er >= 1:
        drop = t - F
        if drop <= 0:
            m = mag << (-drop)
            guard = sticky = 0
        else:
            m = mag >> drop
            guard = (mag >> (drop - 1)) & 1
            sticky = 1 if (mag & ((1 << (drop - 1)) - 1)) else 0
        frac = m & fmask
        if guard and ((frac & 1) or sticky):
            frac += 1
            if frac > fmask:
                frac = 0
                er += 1
        if er >= emax:
            return (s_out << (E + F)) | (emax << F)
        return (s_out << (E + F)) | (er << F) | frac
    sh = (1 - bias - F) - g
    if sh <= 0:
        sig = mag << (-sh)
        guard = sticky = 0
    else:
        sig = mag >> sh
        guard = (mag >> (sh - 1)) & 1
        sticky = 1 if (mag & ((1 << (sh - 1)) - 1)) else 0
    if guard and ((sig & 1) or sticky):
        sig += 1
    if sig == 0:
        return s_out << (E + F)
    if sig >= (1 << F):
        return (s_out << (E + F)) | (1 << F)
    return (s_out << (E + F)) | sig


def float_add_oracle(aw: int, bw: int, exp_bits: int, frac_bits: int) -> int:
    """Expected sum word under the documented contract: exact IEEE specials
    (NaN, infinities, opposite-sign infinities -> NaN, signed zeros, exact
    cancellation -> +0), subnormal operands and gradual-underflow subnormal
    results, round-to-nearest-even mantissa. A zero operand passes the other
    through verbatim. Pure integer arithmetic, so exact."""
    E, F = exp_bits, frac_bits
    emax = (1 << E) - 1
    fmask = (1 << F) - 1
    qnan = (emax << F) | (1 << (F - 1))
    sa, ea, fa = (aw >> (E + F)) & 1, (aw >> F) & emax, aw & fmask
    sb, eb, fb = (bw >> (E + F)) & 1, (bw >> F) & emax, bw & fmask
    a_nan = ea == emax and fa != 0
    b_nan = eb == emax and fb != 0
    a_inf = ea == emax and fa == 0
    b_inf = eb == emax and fb == 0
    if a_nan or b_nan or (a_inf and b_inf and sa != sb):
        return qnan
    if a_inf or b_inf:
        s = sa if a_inf else sb
        return (s << (E + F)) | (emax << F)
    a_zero = ea == 0 and fa == 0                 # true zero, not subnormal
    b_zero = eb == 0 and fb == 0
    if a_zero and b_zero:
        return (sa & sb) << (E + F)
    if a_zero:
        return bw
    if b_zero:
        return aw
    Ma = ((1 << F) if ea else 0) | fa            # implicit bit = (exp != 0)
    Mb = ((1 << F) if eb else 0) | fb
    eea = ea if ea else 1                         # effective (denormal) exponent
    eeb = eb if eb else 1
    pla = (eea << (F + 1)) | Ma
    plb = (eeb << (F + 1)) | Mb
    if pla >= plb:
        sL, ML, eL, MS, eS = sa, Ma, eea, Mb, eeb
    else:
        sL, ML, eL, MS, eS = sb, Mb, eeb, Ma, eea
    d = eL - eS
    # Exact fixed-point value with G extra low bits below the mantissa LSB.
    G = F + 4
    total = (ML << (d + G)) + (MS << G) if sa == sb else (ML << (d + G)) - (MS << G)
    if total == 0:
        return 0  # exact cancellation -> +0
    t = total.bit_length() - 1
    exp_r = eS + (t - G) - F
    if exp_r >= 1:
        lead = F + G
        if t >= lead:
            sh = t - lead
            mant_ext = total >> sh
            sticky_low = 1 if (total & ((1 << sh) - 1)) else 0
        else:
            mant_ext = total << (lead - t)
            sticky_low = 0
        frac = (mant_ext >> G) & fmask
        guard = (mant_ext >> (G - 1)) & 1
        below = (mant_ext & ((1 << (G - 1)) - 1)) or sticky_low
        if guard and ((frac & 1) or below):      # round-to-nearest-even
            frac += 1
            if frac > fmask:
                frac = 0
                exp_r += 1
        if exp_r >= emax:
            return (sL << (E + F)) | (emax << F)
        return (sL << (E + F)) | (exp_r << F) | frac
    # subnormal result: gradual underflow to exponent field 0
    rshift = G + 1 - eS
    if rshift <= 0:
        sig = total << (-rshift)
    else:
        sig = total >> rshift
        guard = (total >> (rshift - 1)) & 1
        sticky = 1 if (total & ((1 << (rshift - 1)) - 1)) else 0
        if guard and ((sig & 1) or sticky):
            sig += 1
    if sig == 0:
        return sL << (E + F)
    if sig >= (1 << F):                           # rounded up to smallest normal
        return (sL << (E + F)) | (1 << F)
    return (sL << (E + F)) | sig


# =============================================================================
# CIRCUIT EVALUATION
# =============================================================================

class BatchedFitnessEvaluator:
    """
    GPU-batched fitness evaluator with per-circuit reporting.
    Tests all circuits comprehensively.
    """

    def __init__(self, device: str = 'cuda', model_path: str = MODEL_PATH, tensors: Dict[str, torch.Tensor] = None):
        self.device = device
        self.model_path = model_path
        self.metadata = load_metadata(model_path)
        self.signal_registry = self.metadata.get('signal_registry', {})
        self.results: List[CircuitResult] = []
        self.category_scores: Dict[str, Tuple[float, int]] = {}
        self.total_tests = 0

        # Get manifest for N-bit support
        if tensors is not None:
            self.manifest = get_manifest(tensors)
        else:
            base_tensors = load_model(model_path)
            self.manifest = get_manifest(base_tensors)
        self.data_bits = self.manifest['data_bits']
        self.addr_bits = self.manifest['addr_bits']

        self._setup_tests()

    def _setup_tests(self):
        """Pre-compute test vectors on device."""
        d = self.device

        # 2-input truth table [4, 2]
        self.tt2 = torch.tensor(
            [[0, 0], [0, 1], [1, 0], [1, 1]],
            device=d, dtype=torch.float32
        )

        # 3-input truth table [8, 3]
        self.tt3 = torch.tensor([
            [0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1],
            [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1]
        ], device=d, dtype=torch.float32)

        # Boolean gate expected outputs
        self.expected = {
            'and': torch.tensor([0, 0, 0, 1], device=d, dtype=torch.float32),
            'or': torch.tensor([0, 1, 1, 1], device=d, dtype=torch.float32),
            'nand': torch.tensor([1, 1, 1, 0], device=d, dtype=torch.float32),
            'nor': torch.tensor([1, 0, 0, 0], device=d, dtype=torch.float32),
            'xor': torch.tensor([0, 1, 1, 0], device=d, dtype=torch.float32),
            'xnor': torch.tensor([1, 0, 0, 1], device=d, dtype=torch.float32),
            'implies': torch.tensor([1, 1, 0, 1], device=d, dtype=torch.float32),
            'biimplies': torch.tensor([1, 0, 0, 1], device=d, dtype=torch.float32),
            'not': torch.tensor([1, 0], device=d, dtype=torch.float32),
            'ha_sum': torch.tensor([0, 1, 1, 0], device=d, dtype=torch.float32),
            'ha_carry': torch.tensor([0, 0, 0, 1], device=d, dtype=torch.float32),
            'fa_sum': torch.tensor([0, 1, 1, 0, 1, 0, 0, 1], device=d, dtype=torch.float32),
            'fa_cout': torch.tensor([0, 0, 0, 1, 0, 1, 1, 1], device=d, dtype=torch.float32),
        }

        # NOT gate inputs
        self.not_inputs = torch.tensor([[0], [1]], device=d, dtype=torch.float32)

        # 8-bit test values
        self.test_8bit = torch.tensor([
            0, 1, 2, 3, 4, 7, 8, 15, 16, 31, 32, 63, 64, 127, 128, 255,
            0b10101010, 0b01010101, 0b11110000, 0b00001111,
            0b11001100, 0b00110011, 0b10000001, 0b01111110
        ], device=d, dtype=torch.long)

        # Bit representations [num_vals, 8]
        self.test_8bit_bits = torch.stack([
            ((self.test_8bit >> (7 - i)) & 1).float() for i in range(8)
        ], dim=1)

        # Comparator test pairs
        comp_tests = [
            (0, 0), (1, 0), (0, 1), (5, 3), (3, 5), (5, 5),
            (255, 0), (0, 255), (128, 127), (127, 128),
            (100, 99), (99, 100), (64, 32), (32, 64),
            (1, 1), (254, 255), (255, 254), (128, 128),
            (0, 128), (128, 0), (64, 64), (192, 192),
            (15, 16), (16, 15), (240, 239), (239, 240),
            (85, 170), (170, 85), (0xAA, 0x55), (0x55, 0xAA),
            (0x0F, 0xF0), (0xF0, 0x0F), (0x33, 0xCC), (0xCC, 0x33),
            (2, 3), (3, 2), (126, 127), (127, 126),
            (129, 128), (128, 129), (200, 199), (199, 200),
            (50, 51), (51, 50), (10, 20), (20, 10),
            (100, 100), (200, 200), (77, 77), (0, 0)
        ]
        self.comp_a = torch.tensor([c[0] for c in comp_tests], device=d, dtype=torch.long)
        self.comp_b = torch.tensor([c[1] for c in comp_tests], device=d, dtype=torch.long)

        # Modular test range
        self.mod_test = torch.arange(256, device=d, dtype=torch.long)

        # 32-bit test values (strategic sampling)
        self.test_32bit = torch.tensor([
            0, 1, 2, 255, 256, 65535, 65536,
            0x7FFFFFFF, 0x80000000, 0xFFFFFFFF,
            0x12345678, 0xDEADBEEF, 0xCAFEBABE,
            1000000, 1000000000, 2147483647,
            0x55555555, 0xAAAAAAAA, 0x0F0F0F0F, 0xF0F0F0F0
        ], device=d, dtype=torch.long)

        # 32-bit comparator test pairs
        comp32_tests = [
            (0, 0), (1, 0), (0, 1), (1000, 999), (999, 1000),
            (0xFFFFFFFF, 0), (0, 0xFFFFFFFF),
            (0x80000000, 0x7FFFFFFF), (0x7FFFFFFF, 0x80000000),
            (1000000, 1000000), (0x12345678, 0x12345678),
            (0xDEADBEEF, 0xCAFEBABE), (0xCAFEBABE, 0xDEADBEEF),
            (256, 255), (255, 256), (65536, 65535), (65535, 65536),
        ]
        self.comp32_a = torch.tensor([c[0] for c in comp32_tests], device=d, dtype=torch.long)
        self.comp32_b = torch.tensor([c[1] for c in comp32_tests], device=d, dtype=torch.long)

    def _record(self, name: str, passed: int, total: int, failures: List[Tuple] = None):
        """Record a circuit test result."""
        self.results.append(CircuitResult(
            name=name,
            passed=passed,
            total=total,
            failures=failures or []
        ))

    # =========================================================================
    # BOOLEAN GATES
    # =========================================================================

    def _test_single_gate(self, pop: Dict, prefix: str, inputs: torch.Tensor,
                          expected: torch.Tensor) -> torch.Tensor:
        """Test single-layer gate (AND, OR, NAND, NOR, IMPLIES)."""
        pop_size = next(iter(pop.values())).shape[0]
        w = pop[f'{prefix}.weight']
        b = pop[f'{prefix}.bias']

        # [num_tests, pop_size]
        out = heaviside(inputs @ w.view(pop_size, -1).T + b.view(pop_size))
        correct = (out == expected.unsqueeze(1)).float().sum(0)

        failures = []
        if pop_size == 1:
            for i, (inp, exp, got) in enumerate(zip(inputs, expected, out[:, 0])):
                if exp.item() != got.item():
                    failures.append((inp.tolist(), exp.item(), got.item()))

        self._record(prefix, int(correct[0].item()), len(expected), failures)
        return correct

    def _test_twolayer_gate(self, pop: Dict, prefix: str, inputs: torch.Tensor,
                            expected: torch.Tensor) -> torch.Tensor:
        """Test two-layer gate (XOR, XNOR, BIIMPLIES)."""
        pop_size = next(iter(pop.values())).shape[0]

        # Layer 1
        w1_n1 = pop[f'{prefix}.layer1.neuron1.weight']
        b1_n1 = pop[f'{prefix}.layer1.neuron1.bias']
        w1_n2 = pop[f'{prefix}.layer1.neuron2.weight']
        b1_n2 = pop[f'{prefix}.layer1.neuron2.bias']

        h1 = heaviside(inputs @ w1_n1.view(pop_size, -1).T + b1_n1.view(pop_size))
        h2 = heaviside(inputs @ w1_n2.view(pop_size, -1).T + b1_n2.view(pop_size))
        hidden = torch.stack([h1, h2], dim=-1)

        # Layer 2
        w2 = pop[f'{prefix}.layer2.weight']
        b2 = pop[f'{prefix}.layer2.bias']
        out = heaviside((hidden * w2.view(pop_size, 2)).sum(-1) + b2.view(pop_size))

        correct = (out == expected.unsqueeze(1)).float().sum(0)

        failures = []
        if pop_size == 1:
            for i, (inp, exp, got) in enumerate(zip(inputs, expected, out[:, 0])):
                if exp.item() != got.item():
                    failures.append((inp.tolist(), exp.item(), got.item()))

        self._record(prefix, int(correct[0].item()), len(expected), failures)
        return correct

    def _test_xor_ornand(self, pop: Dict, prefix: str, inputs: torch.Tensor,
                         expected: torch.Tensor) -> torch.Tensor:
        """Test XOR with or/nand layer naming."""
        pop_size = next(iter(pop.values())).shape[0]

        w_or = pop[f'{prefix}.layer1.or.weight']
        b_or = pop[f'{prefix}.layer1.or.bias']
        w_nand = pop[f'{prefix}.layer1.nand.weight']
        b_nand = pop[f'{prefix}.layer1.nand.bias']

        h_or = heaviside(inputs @ w_or.view(pop_size, -1).T + b_or.view(pop_size))
        h_nand = heaviside(inputs @ w_nand.view(pop_size, -1).T + b_nand.view(pop_size))
        hidden = torch.stack([h_or, h_nand], dim=-1)

        w2 = pop[f'{prefix}.layer2.weight']
        b2 = pop[f'{prefix}.layer2.bias']
        out = heaviside((hidden * w2.view(pop_size, 2)).sum(-1) + b2.view(pop_size))

        correct = (out == expected.unsqueeze(1)).float().sum(0)

        failures = []
        if pop_size == 1:
            for i, (inp, exp, got) in enumerate(zip(inputs, expected, out[:, 0])):
                if exp.item() != got.item():
                    failures.append((inp.tolist(), exp.item(), got.item()))

        self._record(prefix, int(correct[0].item()), len(expected), failures)
        return correct

    def _test_boolean_gates(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test all boolean gates."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== BOOLEAN GATES ===")

        # Single-layer gates
        for gate in ['and', 'or', 'nand', 'nor', 'implies']:
            scores += self._test_single_gate(pop, f'boolean.{gate}', self.tt2, self.expected[gate])
            total += 4
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")

        # NOT gate
        w = pop['boolean.not.weight']
        b = pop['boolean.not.bias']
        out = heaviside(self.not_inputs @ w.view(pop_size, -1).T + b.view(pop_size))
        correct = (out == self.expected['not'].unsqueeze(1)).float().sum(0)
        scores += correct
        total += 2

        failures = []
        if pop_size == 1:
            for inp, exp, got in zip(self.not_inputs, self.expected['not'], out[:, 0]):
                if exp.item() != got.item():
                    failures.append((inp.tolist(), exp.item(), got.item()))
        self._record('boolean.not', int(correct[0].item()), 2, failures)
        if debug:
            r = self.results[-1]
            print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")

        # Two-layer gates
        for gate in ['xnor', 'biimplies']:
            scores += self._test_twolayer_gate(pop, f'boolean.{gate}', self.tt2, self.expected.get(gate, self.expected['xnor']))
            total += 4
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")

        # XOR with neuron1/neuron2 naming (same as xnor/biimplies)
        scores += self._test_twolayer_gate(pop, 'boolean.xor', self.tt2, self.expected['xor'])
        total += 4
        if debug:
            r = self.results[-1]
            print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")

        return scores, total

    # =========================================================================
    # ARITHMETIC - ADDERS
    # =========================================================================

    def _eval_xor(self, pop: Dict, prefix: str, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        """Evaluate XOR gate with or/nand decomposition.

        Args:
            a, b: Tensors of shape [num_tests] or [num_tests, pop_size]

        Returns:
            Tensor of shape [num_tests, pop_size]
        """
        pop_size = next(iter(pop.values())).shape[0]

        # Ensure inputs are [num_tests, pop_size]
        if a.dim() == 1:
            a = a.unsqueeze(1).expand(-1, pop_size)
        if b.dim() == 1:
            b = b.unsqueeze(1).expand(-1, pop_size)

        # inputs: [num_tests, pop_size, 2]
        inputs = torch.stack([a, b], dim=-1)

        w_or = pop[f'{prefix}.layer1.or.weight'].view(pop_size, 2)
        b_or = pop[f'{prefix}.layer1.or.bias'].view(pop_size)
        w_nand = pop[f'{prefix}.layer1.nand.weight'].view(pop_size, 2)
        b_nand = pop[f'{prefix}.layer1.nand.bias'].view(pop_size)

        # [num_tests, pop_size]
        h_or = heaviside((inputs * w_or).sum(-1) + b_or)
        h_nand = heaviside((inputs * w_nand).sum(-1) + b_nand)

        # hidden: [num_tests, pop_size, 2]
        hidden = torch.stack([h_or, h_nand], dim=-1)

        w2 = pop[f'{prefix}.layer2.weight'].view(pop_size, 2)
        b2 = pop[f'{prefix}.layer2.bias'].view(pop_size)
        return heaviside((hidden * w2).sum(-1) + b2)

    def _eval_single_fa(self, pop: Dict, prefix: str,
                        a: torch.Tensor, b: torch.Tensor, cin: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Evaluate single full adder.

        Args:
            a, b, cin: Tensors of shape [num_tests] or [num_tests, pop_size]

        Returns:
            sum_out, cout: Both of shape [num_tests, pop_size]
        """
        pop_size = next(iter(pop.values())).shape[0]

        # Ensure inputs are [num_tests, pop_size]
        if a.dim() == 1:
            a = a.unsqueeze(1).expand(-1, pop_size)
        if b.dim() == 1:
            b = b.unsqueeze(1).expand(-1, pop_size)
        if cin.dim() == 1:
            cin = cin.unsqueeze(1).expand(-1, pop_size)

        # Half adder 1: a XOR b -> [num_tests, pop_size]
        ha1_sum = self._eval_xor(pop, f'{prefix}.ha1.sum', a, b)

        # Half adder 1 carry: a AND b
        ab = torch.stack([a, b], dim=-1)  # [num_tests, pop_size, 2]
        w_c1 = pop[f'{prefix}.ha1.carry.weight'].view(pop_size, 2)
        b_c1 = pop[f'{prefix}.ha1.carry.bias'].view(pop_size)
        ha1_carry = heaviside((ab * w_c1).sum(-1) + b_c1)

        # Half adder 2: ha1_sum XOR cin
        ha2_sum = self._eval_xor(pop, f'{prefix}.ha2.sum', ha1_sum, cin)

        # Half adder 2 carry
        sc = torch.stack([ha1_sum, cin], dim=-1)
        w_c2 = pop[f'{prefix}.ha2.carry.weight'].view(pop_size, 2)
        b_c2 = pop[f'{prefix}.ha2.carry.bias'].view(pop_size)
        ha2_carry = heaviside((sc * w_c2).sum(-1) + b_c2)

        # Carry out: ha1_carry OR ha2_carry
        carries = torch.stack([ha1_carry, ha2_carry], dim=-1)
        w_cout = pop[f'{prefix}.carry_or.weight'].view(pop_size, 2)
        b_cout = pop[f'{prefix}.carry_or.bias'].view(pop_size)
        cout = heaviside((carries * w_cout).sum(-1) + b_cout)

        return ha2_sum, cout

    def _test_halfadder(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test half adder."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== HALF ADDER ===")

        # Sum (XOR)
        scores += self._test_xor_ornand(pop, 'arithmetic.halfadder.sum', self.tt2, self.expected['ha_sum'])
        total += 4
        if debug:
            r = self.results[-1]
            print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")

        # Carry (AND)
        scores += self._test_single_gate(pop, 'arithmetic.halfadder.carry', self.tt2, self.expected['ha_carry'])
        total += 4
        if debug:
            r = self.results[-1]
            print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")

        return scores, total

    def _test_fulladder(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test full adder with all 8 input combinations."""
        pop_size = next(iter(pop.values())).shape[0]

        if debug:
            print("\n=== FULL ADDER ===")

        a = self.tt3[:, 0]
        b = self.tt3[:, 1]
        cin = self.tt3[:, 2]

        sum_out, cout = self._eval_single_fa(pop, 'arithmetic.fulladder', a, b, cin)

        sum_correct = (sum_out == self.expected['fa_sum'].unsqueeze(1)).float().sum(0)
        cout_correct = (cout == self.expected['fa_cout'].unsqueeze(1)).float().sum(0)

        failures_sum = []
        failures_cout = []
        if pop_size == 1:
            for i in range(8):
                if sum_out[i, 0].item() != self.expected['fa_sum'][i].item():
                    failures_sum.append(([a[i].item(), b[i].item(), cin[i].item()],
                                        self.expected['fa_sum'][i].item(), sum_out[i, 0].item()))
                if cout[i, 0].item() != self.expected['fa_cout'][i].item():
                    failures_cout.append(([a[i].item(), b[i].item(), cin[i].item()],
                                         self.expected['fa_cout'][i].item(), cout[i, 0].item()))

        self._record('arithmetic.fulladder.sum', int(sum_correct[0].item()), 8, failures_sum)
        self._record('arithmetic.fulladder.cout', int(cout_correct[0].item()), 8, failures_cout)

        if debug:
            for r in self.results[-2:]:
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")

        return sum_correct + cout_correct, 16

    def _test_ripplecarry(self, pop: Dict, bits: int, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test N-bit ripple carry adder."""
        pop_size = next(iter(pop.values())).shape[0]

        if debug:
            print(f"\n=== RIPPLE CARRY {bits}-BIT ===")

        prefix = f'arithmetic.ripplecarry{bits}bit'
        max_val = 1 << bits
        num_tests = min(max_val * max_val, 65536)

        if bits <= 4:
            # Exhaustive for small widths
            test_a = torch.arange(max_val, device=self.device)
            test_b = torch.arange(max_val, device=self.device)
            a_vals, b_vals = torch.meshgrid(test_a, test_b, indexing='ij')
            a_vals = a_vals.flatten()
            b_vals = b_vals.flatten()
        elif bits == 8:
            # Strategic sampling for 8-bit
            edge_vals = [0, 1, 2, 127, 128, 254, 255]
            pairs = [(a, b) for a in edge_vals for b in edge_vals]
            for i in range(0, 256, 16):
                pairs.append((i, 255 - i))
            pairs = list(set(pairs))
            a_vals = torch.tensor([p[0] for p in pairs], device=self.device)
            b_vals = torch.tensor([p[1] for p in pairs], device=self.device)
            num_tests = len(pairs)
        else:
            # Width-scaled edges and bit patterns so the high full adders see
            # ones and carries, not just the low byte.
            half = 1 << (bits - 1)
            top = max_val - 1
            alt_a = int('AA' * (bits // 8), 16)
            alt_5 = int('55' * (bits // 8), 16)
            edge_vals = [0, 1, 2, 255, 256, half - 1, half, top - 1, top, alt_a, alt_5]
            if bits == 32:
                edge_vals += [0xFFFF, 0x10000, 0xDEADBEEF, 0x12345678, 1_000_000_000]
            pairs = [(a, b) for a in edge_vals for b in edge_vals]
            step = max_val // 16
            for i in range(0, 16):
                v = i * step
                pairs.append((v, top - v))
            pairs = list(set(pairs))
            a_vals = torch.tensor([p[0] for p in pairs], device=self.device)
            b_vals = torch.tensor([p[1] for p in pairs], device=self.device)
            num_tests = len(pairs)

        # Convert to bits [num_tests, bits]
        a_bits = torch.stack([((a_vals >> (bits - 1 - i)) & 1).float() for i in range(bits)], dim=1)
        b_bits = torch.stack([((b_vals >> (bits - 1 - i)) & 1).float() for i in range(bits)], dim=1)

        # Evaluate ripple carry
        carry = torch.zeros(len(a_vals), pop_size, device=self.device)
        sum_bits = []

        for bit in range(bits):
            bit_idx = bits - 1 - bit  # LSB first
            s, carry = self._eval_single_fa(
                pop, f'{prefix}.fa{bit}',
                a_bits[:, bit_idx].unsqueeze(1).expand(-1, pop_size),
                b_bits[:, bit_idx].unsqueeze(1).expand(-1, pop_size),
                carry
            )
            sum_bits.append(s)

        # Reconstruct result. float64 keeps 32-bit values exact; float32
        # would round above 2^24 and corrupt the comparison.
        sum_bits = torch.stack(sum_bits[::-1], dim=-1)  # MSB first
        result = torch.zeros(len(a_vals), pop_size, device=self.device, dtype=torch.float64)
        for i in range(bits):
            result += sum_bits[:, :, i].double() * (1 << (bits - 1 - i))

        # Expected
        expected = ((a_vals + b_vals) & (max_val - 1)).unsqueeze(1).expand(-1, pop_size).double()
        correct = (result == expected).float().sum(0)

        failures = []
        if pop_size == 1:
            for i in range(min(len(a_vals), 100)):
                if result[i, 0].item() != expected[i, 0].item():
                    failures.append((
                        [int(a_vals[i].item()), int(b_vals[i].item())],
                        int(expected[i, 0].item()),
                        int(result[i, 0].item())
                    ))

        self._record(prefix, int(correct[0].item()), num_tests, failures)
        if debug:
            r = self.results[-1]
            print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")

        return correct, num_tests

    # =========================================================================
    # 3-OPERAND ADDER
    # =========================================================================

    def _test_add3(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test 3-operand 8-bit adder (A + B + C)."""
        pop_size = next(iter(pop.values())).shape[0]

        if debug:
            print(f"\n=== 3-OPERAND ADDER ===")

        prefix = 'arithmetic.add3_8bit'
        bits = 8

        # Strategic test cases for 3-operand addition
        # Include edge cases and overflow scenarios
        test_cases = []
        # Small values
        for a in [0, 1, 2]:
            for b in [0, 1, 2]:
                for c in [0, 1, 2]:
                    test_cases.append((a, b, c))
        # Edge values
        edge = [0, 1, 127, 128, 254, 255]
        for a in edge:
            for b in edge:
                for c in edge:
                    test_cases.append((a, b, c))
        # Specific multi-operand expression tests
        test_cases.extend([
            (15, 27, 33),   # Example from roadmap: 15 + 27 + 33 = 75
            (100, 100, 55), # = 255 (exact fit)
            (100, 100, 56), # = 256 -> 0 (overflow)
            (85, 85, 85),   # = 255 (exact fit)
            (86, 85, 85),   # = 256 -> 0 (overflow)
        ])
        test_cases = list(set(test_cases))

        a_vals = torch.tensor([t[0] for t in test_cases], device=self.device)
        b_vals = torch.tensor([t[1] for t in test_cases], device=self.device)
        c_vals = torch.tensor([t[2] for t in test_cases], device=self.device)
        num_tests = len(test_cases)

        # Convert to bits [num_tests, bits] MSB-first
        a_bits = torch.stack([((a_vals >> (bits - 1 - i)) & 1).float() for i in range(bits)], dim=1)
        b_bits = torch.stack([((b_vals >> (bits - 1 - i)) & 1).float() for i in range(bits)], dim=1)
        c_bits = torch.stack([((c_vals >> (bits - 1 - i)) & 1).float() for i in range(bits)], dim=1)

        # Stage 1: A + B
        carry1 = torch.zeros(num_tests, pop_size, device=self.device)
        stage1_bits = []
        for bit in range(bits):
            bit_idx = bits - 1 - bit  # LSB first
            s, carry1 = self._eval_single_fa(
                pop, f'{prefix}.stage1.fa{bit}',
                a_bits[:, bit_idx].unsqueeze(1).expand(-1, pop_size),
                b_bits[:, bit_idx].unsqueeze(1).expand(-1, pop_size),
                carry1
            )
            stage1_bits.append(s)

        # Stage 2: stage1_result + C
        carry2 = torch.zeros(num_tests, pop_size, device=self.device)
        result_bits = []
        for bit in range(bits):
            bit_idx = bits - 1 - bit  # LSB first
            s, carry2 = self._eval_single_fa(
                pop, f'{prefix}.stage2.fa{bit}',
                stage1_bits[bit],  # Already [num_tests, pop_size]
                c_bits[:, bit_idx].unsqueeze(1).expand(-1, pop_size),
                carry2
            )
            result_bits.append(s)

        # Reconstruct result (bits are in LSB-first order, need to reverse for MSB-first)
        result_bits = torch.stack(result_bits[::-1], dim=-1)  # MSB first
        result = torch.zeros(num_tests, pop_size, device=self.device)
        for i in range(bits):
            result += result_bits[:, :, i] * (1 << (bits - 1 - i))

        # Expected (8-bit wrap)
        expected = ((a_vals + b_vals + c_vals) & 0xFF).unsqueeze(1).expand(-1, pop_size).float()
        correct = (result == expected).float().sum(0)

        failures = []
        if pop_size == 1:
            for i in range(min(num_tests, 100)):
                if result[i, 0].item() != expected[i, 0].item():
                    failures.append((
                        [int(a_vals[i].item()), int(b_vals[i].item()), int(c_vals[i].item())],
                        int(expected[i, 0].item()),
                        int(result[i, 0].item())
                    ))

        self._record(prefix, int(correct[0].item()), num_tests, failures)
        if debug:
            r = self.results[-1]
            print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
            if failures:
                for inp, exp, got in failures[:5]:
                    print(f"    FAIL: {inp[0]} + {inp[1]} + {inp[2]} = {exp}, got {got}")

        return correct, num_tests

    # =========================================================================
    # ORDER OF OPERATIONS (A + B × C)
    # =========================================================================

    def _test_expr_add_mul(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test A + B × C expression circuit (order of operations)."""
        pop_size = next(iter(pop.values())).shape[0]

        if debug:
            print(f"\n=== ORDER OF OPERATIONS (A + B × C) ===")

        prefix = 'arithmetic.expr_add_mul'
        bits = 8

        # Test cases for order of operations
        test_cases = []

        # Specific examples from roadmap
        test_cases.extend([
            (5, 3, 2),    # 5 + 3 × 2 = 5 + 6 = 11
            (10, 4, 3),   # 10 + 4 × 3 = 10 + 12 = 22
            (1, 10, 10),  # 1 + 10 × 10 = 1 + 100 = 101
            (0, 15, 17),  # 0 + 15 × 17 = 255
            (1, 15, 17),  # 1 + 15 × 17 = 256 -> 0 (overflow)
            (100, 5, 5),  # 100 + 5 × 5 = 100 + 25 = 125
        ])

        # Edge cases
        test_cases.extend([
            (0, 0, 0),    # 0 + 0 × 0 = 0
            (255, 0, 0),  # 255 + 0 × 0 = 255
            (0, 255, 1),  # 0 + 255 × 1 = 255
            (0, 1, 255),  # 0 + 1 × 255 = 255
            (1, 1, 1),    # 1 + 1 × 1 = 2
            (0, 16, 16),  # 0 + 16 × 16 = 256 -> 0 (overflow)
        ])

        # Systematic small values
        for a in [0, 1, 5, 10]:
            for b in [0, 1, 2, 3]:
                for c in [0, 1, 2, 3]:
                    test_cases.append((a, b, c))

        # Remove duplicates
        test_cases = list(set(test_cases))

        a_vals = torch.tensor([t[0] for t in test_cases], device=self.device)
        b_vals = torch.tensor([t[1] for t in test_cases], device=self.device)
        c_vals = torch.tensor([t[2] for t in test_cases], device=self.device)
        num_tests = len(test_cases)

        # Convert to bits [num_tests, bits] MSB-first
        a_bits = torch.stack([((a_vals >> (bits - 1 - i)) & 1).float() for i in range(bits)], dim=1)
        b_bits = torch.stack([((b_vals >> (bits - 1 - i)) & 1).float() for i in range(bits)], dim=1)
        c_bits = torch.stack([((c_vals >> (bits - 1 - i)) & 1).float() for i in range(bits)], dim=1)

        # Evaluate mask stage: mask[stage][bit] = B[bit] AND C[stage]
        # In the circuit: mask.s[stage].b[bit] operates on positional bits
        # stage 0 = LSB of C (c_bits[:, 7]), stage 7 = MSB of C (c_bits[:, 0])
        # bit 0 = LSB of B (b_bits[:, 7]), bit 7 = MSB of B (b_bits[:, 0])
        masks = torch.zeros(8, num_tests, pop_size, 8, device=self.device)  # [stage, tests, pop, bits]
        for stage in range(8):
            c_stage_bit = c_bits[:, 7 - stage].unsqueeze(1).expand(-1, pop_size)  # C[stage]
            for bit in range(8):
                b_bit_val = b_bits[:, 7 - bit].unsqueeze(1).expand(-1, pop_size)  # B[bit]
                # AND gate
                w = pop.get(f'{prefix}.mul.mask.s{stage}.b{bit}.weight')
                bias = pop.get(f'{prefix}.mul.mask.s{stage}.b{bit}.bias')
                if w is not None and bias is not None:
                    w = w.squeeze(-1)  # [pop]
                    b_tensor = bias.squeeze(-1)  # [pop]
                    # Properly broadcast for batch evaluation
                    inp = torch.stack([b_bit_val, c_stage_bit], dim=-1)  # [tests, pop, 2]
                    out = heaviside(torch.einsum('tpi,pi->tp', inp, w) + b_tensor)
                    masks[stage, :, :, bit] = out

        # Accumulator stages
        # acc[0] = mask[0] (no shift)
        # acc[1] = acc[0] + (mask[1] << 1)
        # ...
        # acc[7] = acc[6] + (mask[7] << 7)
        acc = masks[0].clone()  # [tests, pop, 8] - start with mask[0]

        for stage in range(1, 8):
            # Create shifted mask: (mask[stage] << stage)
            # Shift left by 'stage' positions: bits 0..stage-1 become 0, bit k becomes mask[stage][k-stage]
            shifted_mask = torch.zeros(num_tests, pop_size, 8, device=self.device)
            for bit in range(8):
                if bit >= stage:
                    shifted_mask[:, :, bit] = masks[stage, :, :, bit - stage]
                # else: remains 0

            # Add acc + shifted_mask using full adders
            carry = torch.zeros(num_tests, pop_size, device=self.device)
            new_acc = torch.zeros(num_tests, pop_size, 8, device=self.device)
            for bit in range(8):
                s, carry = self._eval_single_fa(
                    pop, f'{prefix}.mul.acc.s{stage}.fa{bit}',
                    acc[:, :, bit],
                    shifted_mask[:, :, bit],
                    carry
                )
                new_acc[:, :, bit] = s
            acc = new_acc

        # Final add stage: A + acc (multiplication result)
        carry = torch.zeros(num_tests, pop_size, device=self.device)
        result_bits = []
        for bit in range(8):
            a_bit_val = a_bits[:, 7 - bit].unsqueeze(1).expand(-1, pop_size)
            s, carry = self._eval_single_fa(
                pop, f'{prefix}.add.fa{bit}',
                a_bit_val,
                acc[:, :, bit],
                carry
            )
            result_bits.append(s)

        # Reconstruct result
        result_bits = torch.stack(result_bits[::-1], dim=-1)  # MSB first
        result = torch.zeros(num_tests, pop_size, device=self.device)
        for i in range(bits):
            result += result_bits[:, :, i] * (1 << (bits - 1 - i))

        # Expected: A + (B × C), with 8-bit wrap
        expected = ((a_vals + b_vals * c_vals) & 0xFF).unsqueeze(1).expand(-1, pop_size).float()
        correct = (result == expected).float().sum(0)

        failures = []
        if pop_size == 1:
            for i in range(min(num_tests, 100)):
                if result[i, 0].item() != expected[i, 0].item():
                    failures.append((
                        [int(a_vals[i].item()), int(b_vals[i].item()), int(c_vals[i].item())],
                        int(expected[i, 0].item()),
                        int(result[i, 0].item())
                    ))

        self._record(prefix, int(correct[0].item()), num_tests, failures)
        if debug:
            r = self.results[-1]
            print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
            if failures:
                for inp, exp, got in failures[:5]:
                    print(f"    FAIL: {inp[0]} + {inp[1]} × {inp[2]} = {exp}, got {got}")

        return correct, num_tests

    # =========================================================================
    # COMPARATORS
    # =========================================================================

    def _eval_bit_cascade_compare(
        self,
        pop: Dict,
        cmp_prefix: str,
        out_gt: str,
        out_lt: str,
        out_ge: str,
        out_le: str,
        out_eq: str,
        bits: int,
        a_bits_2d: torch.Tensor,
        b_bits_2d: torch.Tensor,
    ) -> Dict[str, torch.Tensor]:
        """Walk the ternary bit-cascade comparator generated by
        build.add_bit_cascade_compare. Returns a dict with gt/lt/ge/le/eq each of
        shape [num_tests, pop_size]. a_bits_2d, b_bits_2d are [num_tests, bits]
        MSB-first.
        """
        pop_size = next(iter(pop.values())).shape[0]

        # Per-bit gt, lt, eq
        gt_b: List[torch.Tensor] = []
        lt_b: List[torch.Tensor] = []
        eq_b: List[torch.Tensor] = []
        for i in range(bits):
            a_i = a_bits_2d[:, i].unsqueeze(1).expand(-1, pop_size)
            b_i = b_bits_2d[:, i].unsqueeze(1).expand(-1, pop_size)
            ab = torch.stack([a_i, b_i], dim=-1)

            w = pop[f'{cmp_prefix}.bit{i}.gt.weight'].view(pop_size, 2)
            bb = pop[f'{cmp_prefix}.bit{i}.gt.bias'].view(pop_size)
            gt_b.append(heaviside((ab * w).sum(-1) + bb))

            w = pop[f'{cmp_prefix}.bit{i}.lt.weight'].view(pop_size, 2)
            bb = pop[f'{cmp_prefix}.bit{i}.lt.bias'].view(pop_size)
            lt_b.append(heaviside((ab * w).sum(-1) + bb))

            w = pop[f'{cmp_prefix}.bit{i}.eq.layer1.and.weight'].view(pop_size, 2)
            bb = pop[f'{cmp_prefix}.bit{i}.eq.layer1.and.bias'].view(pop_size)
            h_and = heaviside((ab * w).sum(-1) + bb)
            w = pop[f'{cmp_prefix}.bit{i}.eq.layer1.nor.weight'].view(pop_size, 2)
            bb = pop[f'{cmp_prefix}.bit{i}.eq.layer1.nor.bias'].view(pop_size)
            h_nor = heaviside((ab * w).sum(-1) + bb)
            hidden = torch.stack([h_and, h_nor], dim=-1)
            w = pop[f'{cmp_prefix}.bit{i}.eq.weight'].view(pop_size, 2)
            bb = pop[f'{cmp_prefix}.bit{i}.eq.bias'].view(pop_size)
            eq_b.append(heaviside((hidden * w).sum(-1) + bb))

        # eq_prefix[i] = AND of eq[0..i-1]
        eq_pref: List[Optional[torch.Tensor]] = [None]
        for i in range(1, bits):
            eq_stack = torch.stack(eq_b[:i], dim=-1)
            w = pop[f'{cmp_prefix}.cascade.eq_prefix.bit{i}.weight'].view(pop_size, i)
            bb = pop[f'{cmp_prefix}.cascade.eq_prefix.bit{i}.bias'].view(pop_size)
            eq_pref.append(heaviside((eq_stack * w).sum(-1) + bb))

        # cascade gt[i], lt[i] = eq_prefix[i] AND gt_b[i] / lt_b[i]
        casc_gt = [gt_b[0]]
        casc_lt = [lt_b[0]]
        for i in range(1, bits):
            inp = torch.stack([eq_pref[i], gt_b[i]], dim=-1)
            w = pop[f'{cmp_prefix}.cascade.gt.bit{i}.weight'].view(pop_size, 2)
            bb = pop[f'{cmp_prefix}.cascade.gt.bit{i}.bias'].view(pop_size)
            casc_gt.append(heaviside((inp * w).sum(-1) + bb))
            inp = torch.stack([eq_pref[i], lt_b[i]], dim=-1)
            w = pop[f'{cmp_prefix}.cascade.lt.bit{i}.weight'].view(pop_size, 2)
            bb = pop[f'{cmp_prefix}.cascade.lt.bit{i}.bias'].view(pop_size)
            casc_lt.append(heaviside((inp * w).sum(-1) + bb))

        # Final OR for GT / LT
        gt_stack = torch.stack(casc_gt, dim=-1)
        w = pop[f'{out_gt}.weight'].view(pop_size, bits)
        bb = pop[f'{out_gt}.bias'].view(pop_size)
        final_gt = heaviside((gt_stack * w).sum(-1) + bb)

        lt_stack = torch.stack(casc_lt, dim=-1)
        w = pop[f'{out_lt}.weight'].view(pop_size, bits)
        bb = pop[f'{out_lt}.bias'].view(pop_size)
        final_lt = heaviside((lt_stack * w).sum(-1) + bb)

        # Final AND for EQ
        eq_stack = torch.stack(eq_b, dim=-1)
        w = pop[f'{out_eq}.weight'].view(pop_size, bits)
        bb = pop[f'{out_eq}.bias'].view(pop_size)
        final_eq = heaviside((eq_stack * w).sum(-1) + bb)

        # GE = NOT(LT) buffer pair, LE = NOT(GT) buffer pair
        w = pop[f'{out_ge}.not_lt.weight'].view(pop_size)
        bb = pop[f'{out_ge}.not_lt.bias'].view(pop_size)
        not_lt = heaviside(final_lt * w + bb)
        w = pop[f'{out_ge}.weight'].view(pop_size)
        bb = pop[f'{out_ge}.bias'].view(pop_size)
        final_ge = heaviside(not_lt * w + bb)

        w = pop[f'{out_le}.not_gt.weight'].view(pop_size)
        bb = pop[f'{out_le}.not_gt.bias'].view(pop_size)
        not_gt = heaviside(final_gt * w + bb)
        w = pop[f'{out_le}.weight'].view(pop_size)
        bb = pop[f'{out_le}.bias'].view(pop_size)
        final_le = heaviside(not_gt * w + bb)

        return {
            "gt": final_gt, "lt": final_lt, "eq": final_eq,
            "ge": final_ge, "le": final_le,
        }

    def _test_comparators(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test 8-bit comparators (bit-cascade)."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== COMPARATORS (8-bit bit-cascade) ===")

        bits = 8
        a_bits = torch.stack([((self.comp_a >> (bits - 1 - i)) & 1).float() for i in range(bits)], dim=1)
        b_bits = torch.stack([((self.comp_b >> (bits - 1 - i)) & 1).float() for i in range(bits)], dim=1)
        try:
            outs = self._eval_bit_cascade_compare(
                pop,
                f"arithmetic.cmp{bits}bit",
                f"arithmetic.greaterthan{bits}bit",
                f"arithmetic.lessthan{bits}bit",
                f"arithmetic.greaterorequal{bits}bit",
                f"arithmetic.lessorequal{bits}bit",
                f"arithmetic.equality{bits}bit",
                bits,
                a_bits,
                b_bits,
            )
        except KeyError:
            return scores, total

        for kind, op in [
            ("gt", lambda a, b: a > b),
            ("lt", lambda a, b: a < b),
            ("ge", lambda a, b: a >= b),
            ("le", lambda a, b: a <= b),
            ("eq", lambda a, b: a == b),
        ]:
            expected = torch.tensor(
                [1.0 if op(a.item(), b.item()) else 0.0 for a, b in zip(self.comp_a, self.comp_b)],
                device=self.device,
            )
            out = outs[kind]
            correct = (out == expected.unsqueeze(1)).float().sum(0)
            scores += correct
            total += len(self.comp_a)
            name_map = {
                "gt": f"arithmetic.greaterthan{bits}bit",
                "lt": f"arithmetic.lessthan{bits}bit",
                "ge": f"arithmetic.greaterorequal{bits}bit",
                "le": f"arithmetic.lessorequal{bits}bit",
                "eq": f"arithmetic.equality{bits}bit",
            }
            failures = []
            if pop_size == 1:
                for i in range(len(self.comp_a)):
                    if out[i, 0].item() != expected[i].item():
                        failures.append((
                            [int(self.comp_a[i].item()), int(self.comp_b[i].item())],
                            expected[i].item(),
                            out[i, 0].item(),
                        ))
            self._record(name_map[kind], int(correct[0].item()), len(self.comp_a), failures)
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        return scores, total

    def _test_comparators_nbits(self, pop: Dict, bits: int, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test N-bit comparator circuits (GT, LT, GE, LE, EQ) via bit-cascade."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print(f"\n=== {bits}-BIT COMPARATORS (bit-cascade) ===")

        if bits == 32:
            comp_a = self.comp32_a
            comp_b = self.comp32_b
        elif bits == 16:
            comp_a = self.comp_a.clamp(0, 65535)
            comp_b = self.comp_b.clamp(0, 65535)
        else:
            comp_a = self.comp_a
            comp_b = self.comp_b

        num_tests = len(comp_a)
        a_bits = torch.stack([((comp_a >> (bits - 1 - i)) & 1).float() for i in range(bits)], dim=1)
        b_bits = torch.stack([((comp_b >> (bits - 1 - i)) & 1).float() for i in range(bits)], dim=1)

        try:
            outs = self._eval_bit_cascade_compare(
                pop,
                f"arithmetic.cmp{bits}bit",
                f"arithmetic.greaterthan{bits}bit",
                f"arithmetic.lessthan{bits}bit",
                f"arithmetic.greaterorequal{bits}bit",
                f"arithmetic.lessorequal{bits}bit",
                f"arithmetic.equality{bits}bit",
                bits,
                a_bits,
                b_bits,
            )
        except KeyError:
            return scores, total

        for kind, op in [
            ("gt", lambda a, b: a > b),
            ("lt", lambda a, b: a < b),
            ("ge", lambda a, b: a >= b),
            ("le", lambda a, b: a <= b),
            ("eq", lambda a, b: a == b),
        ]:
            expected = torch.tensor(
                [1.0 if op(a.item(), b.item()) else 0.0 for a, b in zip(comp_a, comp_b)],
                device=self.device,
            )
            out = outs[kind]
            correct = (out == expected.unsqueeze(1)).float().sum(0)
            scores += correct
            total += num_tests
            name_map = {
                "gt": f"arithmetic.greaterthan{bits}bit",
                "lt": f"arithmetic.lessthan{bits}bit",
                "ge": f"arithmetic.greaterorequal{bits}bit",
                "le": f"arithmetic.lessorequal{bits}bit",
                "eq": f"arithmetic.equality{bits}bit",
            }
            failures = []
            if pop_size == 1:
                for i in range(num_tests):
                    if out[i, 0].item() != expected[i].item():
                        failures.append((
                            [int(comp_a[i].item()), int(comp_b[i].item())],
                            expected[i].item(),
                            out[i, 0].item(),
                        ))
            self._record(name_map[kind], int(correct[0].item()), num_tests, failures)
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")

        return scores, total

    def _test_subtractor_nbits(self, pop: Dict, bits: int, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test N-bit subtractor circuit (A - B)."""
        pop_size = next(iter(pop.values())).shape[0]

        if debug:
            print(f"\n=== {bits}-BIT SUBTRACTOR ===")

        prefix = f'arithmetic.sub{bits}bit'
        max_val = 1 << bits

        if bits == 32:
            test_pairs = [
                (1000, 500), (5000, 3000), (1000000, 500000),
                (0xFFFFFFFF, 1), (0x80000000, 1), (100, 100),
                (0, 0), (1, 0), (0, 1), (256, 255),
                (0xDEADBEEF, 0xCAFEBABE), (1000000000, 999999999),
            ]
        else:
            test_pairs = [(a, b) for a in [0, 1, 127, 128, 255] for b in [0, 1, 127, 128, 255]]

        a_vals = torch.tensor([p[0] for p in test_pairs], device=self.device, dtype=torch.long)
        b_vals = torch.tensor([p[1] for p in test_pairs], device=self.device, dtype=torch.long)
        num_tests = len(test_pairs)

        a_bits = torch.stack([((a_vals >> (bits - 1 - i)) & 1).float() for i in range(bits)], dim=1)
        b_bits = torch.stack([((b_vals >> (bits - 1 - i)) & 1).float() for i in range(bits)], dim=1)

        not_b_bits = torch.zeros_like(b_bits)
        for bit in range(bits):
            w = pop[f'{prefix}.not_b.bit{bit}.weight'].view(pop_size, -1)
            b = pop[f'{prefix}.not_b.bit{bit}.bias'].view(pop_size)
            not_b_bits[:, bit] = heaviside(b_bits[:, bit:bit+1] @ w.T + b)[:, 0]

        carry = torch.ones(num_tests, pop_size, device=self.device)
        sum_bits = []

        for bit in range(bits):
            bit_idx = bits - 1 - bit
            s, carry = self._eval_single_fa(
                pop, f'{prefix}.fa{bit}',
                a_bits[:, bit_idx].unsqueeze(1).expand(-1, pop_size),
                not_b_bits[:, bit_idx].unsqueeze(1).expand(-1, pop_size),
                carry
            )
            sum_bits.append(s)

        sum_bits = torch.stack(sum_bits[::-1], dim=-1)
        result = torch.zeros(num_tests, pop_size, device=self.device, dtype=torch.float64)
        for i in range(bits):
            result += sum_bits[:, :, i].double() * (1 << (bits - 1 - i))

        expected = ((a_vals - b_vals) & (max_val - 1)).unsqueeze(1).expand(-1, pop_size).double()
        correct = (result == expected).float().sum(0)

        failures = []
        if pop_size == 1:
            for i in range(min(num_tests, 20)):
                if result[i, 0].item() != expected[i, 0].item():
                    failures.append((
                        [int(a_vals[i].item()), int(b_vals[i].item())],
                        int(expected[i, 0].item()),
                        int(result[i, 0].item())
                    ))

        self._record(prefix, int(correct[0].item()), num_tests, failures)
        if debug:
            r = self.results[-1]
            print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")

        return correct, num_tests

    def _test_bitwise_nbits(self, pop: Dict, bits: int, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test N-bit bitwise operations (AND, OR, XOR, NOT)."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print(f"\n=== {bits}-BIT BITWISE OPS ===")

        if bits == 32:
            test_pairs = [
                (0xAAAAAAAA, 0x55555555), (0xFFFFFFFF, 0x00000000),
                (0x12345678, 0x87654321), (0xDEADBEEF, 0xCAFEBABE),
                (0x0F0F0F0F, 0xF0F0F0F0), (0, 0), (0xFFFFFFFF, 0xFFFFFFFF),
            ]
        else:
            test_pairs = [(0xAA, 0x55), (0xFF, 0x00), (0x0F, 0xF0)]

        a_vals = torch.tensor([p[0] for p in test_pairs], device=self.device, dtype=torch.long)
        b_vals = torch.tensor([p[1] for p in test_pairs], device=self.device, dtype=torch.long)
        num_tests = len(test_pairs)

        ops = [
            ('and', lambda a, b: a & b),
            ('or', lambda a, b: a | b),
            ('xor', lambda a, b: a ^ b),
        ]

        for op_name, op_fn in ops:
            try:
                result_bits = []
                for bit in range(bits):
                    a_bit = ((a_vals >> (bits - 1 - bit)) & 1).float().unsqueeze(1).expand(-1, pop_size)
                    b_bit = ((b_vals >> (bits - 1 - bit)) & 1).float().unsqueeze(1).expand(-1, pop_size)
                    inp = torch.stack([a_bit, b_bit], dim=-1)  # [tests, pop, 2]

                    if op_name == 'xor':
                        prefix = f'alu.alu{bits}bit.{op_name}.bit{bit}'
                        w_or = pop[f'{prefix}.layer1.or.weight'].view(pop_size, 2)
                        b_or = pop[f'{prefix}.layer1.or.bias'].view(pop_size)
                        w_nand = pop[f'{prefix}.layer1.nand.weight'].view(pop_size, 2)
                        b_nand = pop[f'{prefix}.layer1.nand.bias'].view(pop_size)
                        h_or = heaviside((inp * w_or).sum(-1) + b_or)
                        h_nand = heaviside((inp * w_nand).sum(-1) + b_nand)
                        hidden = torch.stack([h_or, h_nand], dim=-1)
                        w2 = pop[f'{prefix}.layer2.weight'].view(pop_size, 2)
                        b2 = pop[f'{prefix}.layer2.bias'].view(pop_size)
                        out = heaviside((hidden * w2).sum(-1) + b2)
                    else:
                        w = pop[f'alu.alu{bits}bit.{op_name}.bit{bit}.weight'].view(pop_size, 2)
                        b = pop[f'alu.alu{bits}bit.{op_name}.bit{bit}.bias'].view(pop_size)
                        out = heaviside((inp * w).sum(-1) + b)

                    result_bits.append(out)  # [tests, pop]

                results = torch.zeros(num_tests, pop_size, device=self.device, dtype=torch.float64)
                for i in range(bits):
                    results += result_bits[i].double() * (1 << (bits - 1 - i))
                expected = torch.tensor([op_fn(a.item(), b.item()) for a, b in zip(a_vals, b_vals)],
                                        device=self.device, dtype=torch.float64).unsqueeze(1)

                correct = (results == expected).float().sum(0)  # [pop]
                self._record(f'alu.alu{bits}bit.{op_name}', int(correct[0].item()), num_tests, [])
                if debug:
                    r = self.results[-1]
                    print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
                scores += correct
                total += num_tests
            except KeyError as e:
                if debug:
                    print(f"  alu.alu{bits}bit.{op_name}: SKIP (missing {e})")

        try:
            test_vals = a_vals
            result_bits = []
            for bit in range(bits):
                a_bit = ((test_vals >> (bits - 1 - bit)) & 1).float().unsqueeze(1).expand(-1, pop_size)
                w = pop[f'alu.alu{bits}bit.not.bit{bit}.weight'].view(pop_size)
                b = pop[f'alu.alu{bits}bit.not.bit{bit}.bias'].view(pop_size)
                result_bits.append(heaviside(a_bit * w + b))

            results = torch.zeros(num_tests, pop_size, device=self.device, dtype=torch.float64)
            for i in range(bits):
                results += result_bits[i].double() * (1 << (bits - 1 - i))
            expected = torch.tensor([(~a.item()) & ((1 << bits) - 1) for a in test_vals],
                                    device=self.device, dtype=torch.float64).unsqueeze(1)

            correct = (results == expected).float().sum(0)
            self._record(f'alu.alu{bits}bit.not', int(correct[0].item()), num_tests, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
            scores += correct
            total += num_tests
        except KeyError as e:
            if debug:
                print(f"  alu.alu{bits}bit.not: SKIP (missing {e})")

        return scores, total

    def _test_shifts_nbits(self, pop: Dict, bits: int, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test N-bit shift operations (SHL, SHR)."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print(f"\n=== {bits}-BIT SHIFTS ===")

        if bits == 32:
            test_vals = [0x12345678, 0x80000001, 0x00000001, 0xFFFFFFFF, 0x55555555]
        else:
            test_vals = [0x81, 0x55, 0x01, 0xFF, 0xAA]

        a_vals = torch.tensor(test_vals, device=self.device, dtype=torch.long)
        num_tests = len(test_vals)
        max_val = (1 << bits) - 1

        for op_name, op_fn in [('shl', lambda x: (x << 1) & max_val), ('shr', lambda x: x >> 1)]:
            try:
                result_bits = []
                for bit in range(bits):
                    w = pop[f'alu.alu{bits}bit.{op_name}.bit{bit}.weight'].view(pop_size)
                    b = pop[f'alu.alu{bits}bit.{op_name}.bit{bit}.bias'].view(pop_size)

                    if op_name == 'shl':
                        if bit < bits - 1:
                            src_bit = ((a_vals >> (bits - 2 - bit)) & 1).float()
                        else:
                            src_bit = torch.zeros(num_tests, device=self.device)
                    else:
                        if bit > 0:
                            src_bit = ((a_vals >> (bits - bit)) & 1).float()
                        else:
                            src_bit = torch.zeros(num_tests, device=self.device)

                    out = heaviside(src_bit.unsqueeze(1) * w + b)  # [tests, pop]
                    result_bits.append(out)

                results = torch.zeros(num_tests, pop_size, device=self.device, dtype=torch.float64)
                for i in range(bits):
                    results += result_bits[i].double() * (1 << (bits - 1 - i))
                expected = torch.tensor([op_fn(a.item()) for a in a_vals],
                                        device=self.device, dtype=torch.float64).unsqueeze(1)

                correct = (results == expected).float().sum(0)  # [pop]
                self._record(f'alu.alu{bits}bit.{op_name}', int(correct[0].item()), num_tests, [])
                if debug:
                    r = self.results[-1]
                    print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
                scores += correct
                total += num_tests
            except KeyError as e:
                if debug:
                    print(f"  alu.alu{bits}bit.{op_name}: SKIP (missing {e})")

        return scores, total

    def _test_inc_dec_nbits(self, pop: Dict, bits: int, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test N-bit INC and DEC operations."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print(f"\n=== {bits}-BIT INC/DEC ===")

        if bits == 32:
            test_vals = [0, 1, 0xFFFFFFFF, 0x7FFFFFFF, 0x80000000, 1000000, 0xFFFFFFFE]
        else:
            test_vals = [0, 1, 254, 255, 127, 128]

        a_vals = torch.tensor(test_vals, device=self.device, dtype=torch.long)
        num_tests = len(test_vals)
        max_val = (1 << bits) - 1

        for op_name, op_fn in [('inc', lambda x: (x + 1) & max_val), ('dec', lambda x: (x - 1) & max_val)]:
            try:
                carry = torch.ones(num_tests, pop_size, device=self.device)
                result_bits = []

                for bit in range(bits):
                    a_bit = ((a_vals >> bit) & 1).float().unsqueeze(1).expand(-1, pop_size)

                    prefix = f'alu.alu{bits}bit.{op_name}.bit{bit}'
                    w_or = pop[f'{prefix}.xor.layer1.or.weight'].view(pop_size, 2)
                    b_or = pop[f'{prefix}.xor.layer1.or.bias'].view(pop_size)
                    w_nand = pop[f'{prefix}.xor.layer1.nand.weight'].view(pop_size, 2)
                    b_nand = pop[f'{prefix}.xor.layer1.nand.bias'].view(pop_size)

                    inp = torch.stack([a_bit, carry], dim=-1)  # [tests, pop, 2]
                    h_or = heaviside((inp * w_or).sum(-1) + b_or)
                    h_nand = heaviside((inp * w_nand).sum(-1) + b_nand)

                    w2 = pop[f'{prefix}.xor.layer2.weight'].view(pop_size, 2)
                    b2 = pop[f'{prefix}.xor.layer2.bias'].view(pop_size)
                    hidden = torch.stack([h_or, h_nand], dim=-1)
                    xor_out = heaviside((hidden * w2).sum(-1) + b2)
                    result_bits.append(xor_out)

                    if op_name == 'inc':
                        w_carry = pop[f'{prefix}.carry.weight'].view(pop_size, 2)
                        b_carry = pop[f'{prefix}.carry.bias'].view(pop_size)
                        carry = heaviside((inp * w_carry).sum(-1) + b_carry)
                    else:
                        w_not = pop[f'{prefix}.not_a.weight'].view(pop_size)
                        b_not = pop[f'{prefix}.not_a.bias'].view(pop_size)
                        not_a = heaviside(a_bit * w_not + b_not)
                        w_borrow = pop[f'{prefix}.borrow.weight'].view(pop_size, 2)
                        b_borrow = pop[f'{prefix}.borrow.bias'].view(pop_size)
                        binp = torch.stack([not_a, carry], dim=-1)
                        carry = heaviside((binp * w_borrow).sum(-1) + b_borrow)

                # result_bits[bit] is [tests, pop]; bit index is LSB-first
                results = torch.zeros(num_tests, pop_size, device=self.device, dtype=torch.float64)
                for bit in range(bits):
                    results += result_bits[bit].double() * (1 << bit)
                expected = torch.tensor([op_fn(a.item()) for a in a_vals],
                                        device=self.device, dtype=torch.float64).unsqueeze(1)

                correct = (results == expected).float().sum(0)  # [pop]
                self._record(f'alu.alu{bits}bit.{op_name}', int(correct[0].item()), num_tests, [])
                if debug:
                    r = self.results[-1]
                    print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
                scores += correct
                total += num_tests
            except KeyError as e:
                if debug:
                    print(f"  alu.alu{bits}bit.{op_name}: SKIP (missing {e})")

        return scores, total

    def _test_neg_nbits(self, pop: Dict, bits: int, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test N-bit NEG operation (two's complement negation)."""
        pop_size = next(iter(pop.values())).shape[0]

        if debug:
            print(f"\n=== {bits}-BIT NEG ===")

        if bits == 32:
            test_vals = [0, 1, 0xFFFFFFFF, 0x7FFFFFFF, 0x80000000, 1000, 1000000]
        else:
            test_vals = [0, 1, 127, 128, 255, 100]

        a_vals = torch.tensor(test_vals, device=self.device, dtype=torch.long)
        num_tests = len(test_vals)
        max_val = (1 << bits) - 1

        try:
            not_bits = []
            for bit in range(bits):
                a_bit = ((a_vals >> bit) & 1).float().unsqueeze(1).expand(-1, pop_size)
                w = pop[f'alu.alu{bits}bit.neg.not.bit{bit}.weight'].view(pop_size)
                b = pop[f'alu.alu{bits}bit.neg.not.bit{bit}.bias'].view(pop_size)
                not_bits.append(heaviside(a_bit * w + b))

            carry = torch.ones(num_tests, pop_size, device=self.device)
            result_bits = []

            for bit in range(bits):
                prefix = f'alu.alu{bits}bit.neg.inc.bit{bit}'
                not_bit = not_bits[bit]

                w_or = pop[f'{prefix}.xor.layer1.or.weight'].view(pop_size, 2)
                b_or = pop[f'{prefix}.xor.layer1.or.bias'].view(pop_size)
                w_nand = pop[f'{prefix}.xor.layer1.nand.weight'].view(pop_size, 2)
                b_nand = pop[f'{prefix}.xor.layer1.nand.bias'].view(pop_size)

                inp = torch.stack([not_bit, carry], dim=-1)  # [tests, pop, 2]
                h_or = heaviside((inp * w_or).sum(-1) + b_or)
                h_nand = heaviside((inp * w_nand).sum(-1) + b_nand)

                w2 = pop[f'{prefix}.xor.layer2.weight'].view(pop_size, 2)
                b2 = pop[f'{prefix}.xor.layer2.bias'].view(pop_size)
                hidden = torch.stack([h_or, h_nand], dim=-1)
                xor_out = heaviside((hidden * w2).sum(-1) + b2)
                result_bits.append(xor_out)

                w_carry = pop[f'{prefix}.carry.weight'].view(pop_size, 2)
                b_carry = pop[f'{prefix}.carry.bias'].view(pop_size)
                carry = heaviside((inp * w_carry).sum(-1) + b_carry)

            results = torch.zeros(num_tests, pop_size, device=self.device, dtype=torch.float64)
            for bit in range(bits):
                results += result_bits[bit].double() * (1 << bit)
            expected = torch.tensor([(-a.item()) & max_val for a in a_vals],
                                    device=self.device, dtype=torch.float64).unsqueeze(1)

            correct = (results == expected).float().sum(0)  # [pop]
            self._record(f'alu.alu{bits}bit.neg', int(correct[0].item()), num_tests, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")

            return correct, num_tests
        except KeyError as e:
            if debug:
                print(f"  alu.alu{bits}bit.neg: SKIP (missing {e})")
            return torch.zeros(pop_size, device=self.device), 0

    # =========================================================================
    # THRESHOLD GATES
    # =========================================================================

    def _test_threshold_kofn(self, pop: Dict, k: int, name: str, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test k-of-n threshold gate."""
        pop_size = next(iter(pop.values())).shape[0]
        prefix = f'threshold.{name}'

        # Test all 256 8-bit patterns
        inputs = self.test_8bit_bits if len(self.test_8bit_bits) == 24 else None
        if inputs is None:
            test_vals = torch.arange(256, device=self.device, dtype=torch.long)
            inputs = torch.stack([((test_vals >> (7 - i)) & 1).float() for i in range(8)], dim=1)

        # For k-of-8: output 1 if popcount >= k (for "at least k")
        # For exact naming like "oneoutof8", it's exactly k=1
        popcounts = inputs.sum(dim=1)

        if 'atleast' in name:
            expected = (popcounts >= k).float()
        elif 'atmost' in name or 'minority' in name:
            # minority = popcount <= 3 (less than half of 8)
            expected = (popcounts <= k).float()
        elif 'exactly' in name:
            expected = (popcounts == k).float()
        else:
            # Standard k-of-n (at least k), including majority (>= 5)
            expected = (popcounts >= k).float()

        w = pop[f'{prefix}.weight']
        b = pop[f'{prefix}.bias']
        out = heaviside(inputs @ w.view(pop_size, -1).T + b.view(pop_size))

        correct = (out == expected.unsqueeze(1)).float().sum(0)

        failures = []
        if pop_size == 1:
            for i in range(min(len(inputs), 256)):
                if out[i, 0].item() != expected[i].item():
                    val = int(sum(inputs[i, j].item() * (1 << (7 - j)) for j in range(8)))
                    failures.append((val, expected[i].item(), out[i, 0].item()))

        self._record(prefix, int(correct[0].item()), len(inputs), failures[:10])
        if debug:
            r = self.results[-1]
            print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")

        return correct, len(inputs)

    def _test_threshold_gates(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test all threshold gates."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== THRESHOLD GATES ===")

        # k-of-8 gates
        kofn_gates = [
            (1, 'oneoutof8'), (2, 'twooutof8'), (3, 'threeoutof8'), (4, 'fouroutof8'),
            (5, 'fiveoutof8'), (6, 'sixoutof8'), (7, 'sevenoutof8'), (8, 'alloutof8'),
        ]

        for k, name in kofn_gates:
            try:
                s, t = self._test_threshold_kofn(pop, k, name, debug)
                scores += s
                total += t
            except KeyError:
                pass

        # Special gates
        special = [
            (5, 'majority'), (3, 'minority'),
            (4, 'atleastk_4'), (4, 'atmostk_4'), (4, 'exactlyk_4'),
        ]

        for k, name in special:
            try:
                s, t = self._test_threshold_kofn(pop, k, name, debug)
                scores += s
                total += t
            except KeyError:
                pass

        return scores, total

    # =========================================================================
    # MODULAR ARITHMETIC
    # =========================================================================

    def _test_modular(self, pop: Dict, mod: int, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test modular divisibility circuit.

        Two structures: mod 3/5/6/7/9/10/11/12 use bit-cascade equality
        per multiple of N (`{prefix}.eq.k{k}.*` + final OR at `{prefix}`).
        mod 2/4/8 use a single-layer ternary detector at `{prefix}` directly.
        """
        pop_size = next(iter(pop.values())).shape[0]
        prefix = f'modular.mod{mod}'

        inputs = torch.stack([((self.mod_test >> (7 - i)) & 1).float() for i in range(8)], dim=1)
        expected = ((self.mod_test % mod) == 0).float()
        out = None

        # Bit-cascade equality structure (non-power-of-2 moduli)
        multiples = list(range(0, 256, mod))
        if (multiples
                and f'{prefix}.eq.k{multiples[0]}.all.weight' in pop
                and f'{prefix}.weight' in pop):
            try:
                match_outputs = []
                for k in multiples:
                    per_bit = []
                    for i in range(8):
                        bit_in = inputs[:, i].unsqueeze(1).expand(-1, pop_size)
                        w = pop[f'{prefix}.eq.k{k}.bit{i}.match.weight'].view(pop_size)
                        b = pop[f'{prefix}.eq.k{k}.bit{i}.match.bias'].view(pop_size)
                        per_bit.append(heaviside(bit_in * w + b))
                    stacked = torch.stack(per_bit, dim=-1)
                    w_and = pop[f'{prefix}.eq.k{k}.all.weight'].view(pop_size, 8)
                    b_and = pop[f'{prefix}.eq.k{k}.all.bias'].view(pop_size)
                    match_outputs.append(heaviside((stacked * w_and).sum(-1) + b_and))
                or_in = torch.stack(match_outputs, dim=-1)
                w_or = pop[f'{prefix}.weight'].view(pop_size, len(multiples))
                b_or = pop[f'{prefix}.bias'].view(pop_size)
                out = heaviside((or_in * w_or).sum(-1) + b_or)
            except (KeyError, RuntimeError):
                out = None

        # Single-layer ternary detector (powers of 2)
        if out is None:
            try:
                w = pop[f'{prefix}.weight']
                b = pop[f'{prefix}.bias']
                out = heaviside(inputs @ w.view(pop_size, -1).T + b.view(pop_size))
            except (KeyError, RuntimeError):
                return torch.zeros(pop_size, device=self.device), 0

        correct = (out == expected.unsqueeze(1)).float().sum(0)

        failures = []
        if pop_size == 1:
            for i in range(256):
                if out[i, 0].item() != expected[i].item():
                    failures.append((i, expected[i].item(), out[i, 0].item()))

        self._record(prefix, int(correct[0].item()), 256, failures[:10])
        if debug:
            r = self.results[-1]
            print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")

        return correct, 256

    def _test_modular_all(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test all modular arithmetic circuits."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== MODULAR ARITHMETIC ===")

        for mod in range(2, 13):
            s, t = self._test_modular(pop, mod, debug)
            scores += s
            total += t

        return scores, total

    # =========================================================================
    # PATTERN RECOGNITION
    # =========================================================================

    def _test_pattern(self, pop: Dict, name: str, expected_fn: Callable[[int], float],
                      debug: bool) -> Tuple[torch.Tensor, int]:
        """Test pattern recognition circuit."""
        pop_size = next(iter(pop.values())).shape[0]
        prefix = f'pattern_recognition.{name}'

        test_vals = torch.arange(256, device=self.device, dtype=torch.long)
        inputs = torch.stack([((test_vals >> (7 - i)) & 1).float() for i in range(8)], dim=1)
        expected = torch.tensor([expected_fn(v.item()) for v in test_vals], device=self.device)

        try:
            w = pop[f'{prefix}.weight'].view(pop_size, -1)
            b = pop[f'{prefix}.bias'].view(pop_size)
            out = heaviside(inputs @ w.T + b)
        except KeyError:
            return torch.zeros(pop_size, device=self.device), 0

        correct = (out == expected.unsqueeze(1)).float().sum(0)

        failures = []
        if pop_size == 1:
            for i in range(256):
                if out[i, 0].item() != expected[i].item():
                    failures.append((i, expected[i].item(), out[i, 0].item()))

        self._record(prefix, int(correct[0].item()), 256, failures[:10])
        if debug:
            r = self.results[-1]
            print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")

        return correct, 256

    def _test_patterns(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test pattern recognition circuits."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== PATTERN RECOGNITION ===")

        # Use correct naming: pattern_recognition.allzeros, pattern_recognition.allones
        patterns = [
            ('allzeros', lambda v: 1.0 if v == 0 else 0.0),
            ('allones', lambda v: 1.0 if v == 255 else 0.0),
        ]

        for name, fn in patterns:
            s, t = self._test_pattern(pop, name, fn, debug)
            scores += s
            total += t

        return scores, total

    # =========================================================================
    # ERROR DETECTION
    # =========================================================================

    def _eval_xor_tree_stage(self, pop: Dict, prefix: str, stage: int, idx: int,
                             a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        """Evaluate a single XOR in the parity tree."""
        pop_size = next(iter(pop.values())).shape[0]
        xor_prefix = f'{prefix}.stage{stage}.xor{idx}'

        # Ensure 2D: [256, pop_size]
        if a.dim() == 1:
            a = a.unsqueeze(1).expand(-1, pop_size)
        if b.dim() == 1:
            b = b.unsqueeze(1).expand(-1, pop_size)

        # Layer 1: OR and NAND
        w_or = pop[f'{xor_prefix}.layer1.or.weight'].view(pop_size, 2)
        b_or = pop[f'{xor_prefix}.layer1.or.bias'].view(pop_size)
        w_nand = pop[f'{xor_prefix}.layer1.nand.weight'].view(pop_size, 2)
        b_nand = pop[f'{xor_prefix}.layer1.nand.bias'].view(pop_size)

        inputs = torch.stack([a, b], dim=-1)  # [256, pop_size, 2]
        h_or = heaviside((inputs * w_or).sum(-1) + b_or)
        h_nand = heaviside((inputs * w_nand).sum(-1) + b_nand)

        # Layer 2
        hidden = torch.stack([h_or, h_nand], dim=-1)
        w2 = pop[f'{xor_prefix}.layer2.weight'].view(pop_size, 2)
        b2 = pop[f'{xor_prefix}.layer2.bias'].view(pop_size)
        return heaviside((hidden * w2).sum(-1) + b2)

    def _test_parity_xor_tree(self, pop: Dict, prefix: str, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test parity circuit with XOR tree structure."""
        pop_size = next(iter(pop.values())).shape[0]

        test_vals = torch.arange(256, device=self.device, dtype=torch.long)
        inputs = torch.stack([((test_vals >> (7 - i)) & 1).float() for i in range(8)], dim=1)

        # XOR of all bits: 1 if odd number of 1s
        popcounts = inputs.sum(dim=1)
        xor_result = (popcounts.long() % 2).float()

        try:
            # Stage 1: 4 XORs (pairs of bits)
            s1_out = []
            for i in range(4):
                xor_out = self._eval_xor_tree_stage(pop, prefix, 1, i, inputs[:, i*2], inputs[:, i*2+1])
                s1_out.append(xor_out)

            # Stage 2: 2 XORs
            s2_out = []
            for i in range(2):
                xor_out = self._eval_xor_tree_stage(pop, prefix, 2, i, s1_out[i*2], s1_out[i*2+1])
                s2_out.append(xor_out)

            # Stage 3: 1 XOR
            s3_out = self._eval_xor_tree_stage(pop, prefix, 3, 0, s2_out[0], s2_out[1])

            # Output NOT (for parity checker - inverts the XOR result)
            if f'{prefix}.output.not.weight' in pop:
                w_not = pop[f'{prefix}.output.not.weight'].view(pop_size)
                b_not = pop[f'{prefix}.output.not.bias'].view(pop_size)
                out = heaviside(s3_out * w_not + b_not)
                # Checker outputs 1 if even parity (XOR=0), so expected is inverted xor_result
                expected = 1.0 - xor_result
            else:
                out = s3_out
                expected = xor_result

        except KeyError as e:
            return torch.zeros(pop_size, device=self.device), 0

        correct = (out == expected.unsqueeze(1)).float().sum(0)

        failures = []
        if pop_size == 1:
            for i in range(256):
                if out[i, 0].item() != expected[i].item():
                    failures.append((i, expected[i].item(), out[i, 0].item()))

        self._record(prefix, int(correct[0].item()), 256, failures[:10])
        if debug:
            r = self.results[-1]
            print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")

        return correct, 256

    def _test_error_detection(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test error detection circuits."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== ERROR DETECTION ===")

        # XOR tree parity circuits
        for prefix in ['error_detection.paritychecker8bit', 'error_detection.paritygenerator8bit']:
            s, t = self._test_parity_xor_tree(pop, prefix, debug)
            scores += s
            total += t

        return scores, total

    # =========================================================================
    # COMBINATIONAL LOGIC
    # =========================================================================

    def _test_mux2to1(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test 2-to-1 multiplexer."""
        pop_size = next(iter(pop.values())).shape[0]
        prefix = 'combinational.multiplexer2to1'

        # Inputs: [a, b, sel] -> out = sel ? b : a
        inputs = torch.tensor([
            [0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1],
            [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1],
        ], device=self.device, dtype=torch.float32)
        expected = torch.tensor([0, 0, 0, 1, 1, 0, 1, 1], device=self.device, dtype=torch.float32)

        try:
            w = pop[f'{prefix}.weight']
            b = pop[f'{prefix}.bias']
            out = heaviside(inputs @ w.view(pop_size, -1).T + b.view(pop_size))
        except KeyError:
            return torch.zeros(pop_size, device=self.device), 0

        correct = (out == expected.unsqueeze(1)).float().sum(0)

        failures = []
        if pop_size == 1:
            for i in range(8):
                if out[i, 0].item() != expected[i].item():
                    failures.append((inputs[i].tolist(), expected[i].item(), out[i, 0].item()))

        self._record(prefix, int(correct[0].item()), 8, failures)
        if debug:
            r = self.results[-1]
            print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")

        return correct, 8

    def _test_decoder3to8(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test 3-to-8 decoder."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== DECODER 3-TO-8 ===")

        inputs = torch.tensor([
            [0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1],
            [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1],
        ], device=self.device, dtype=torch.float32)

        for out_idx in range(8):
            prefix = f'combinational.decoder3to8.out{out_idx}'
            expected = torch.zeros(8, device=self.device)
            expected[out_idx] = 1.0

            try:
                w = pop[f'{prefix}.weight']
                b = pop[f'{prefix}.bias']
                out = heaviside(inputs @ w.view(pop_size, -1).T + b.view(pop_size))
            except KeyError:
                continue

            correct = (out == expected.unsqueeze(1)).float().sum(0)
            scores += correct
            total += 8

            failures = []
            if pop_size == 1:
                for i in range(8):
                    if out[i, 0].item() != expected[i].item():
                        failures.append((inputs[i].tolist(), expected[i].item(), out[i, 0].item()))

            self._record(prefix, int(correct[0].item()), 8, failures)
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")

        return scores, total

    def _test_combinational(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test combinational logic circuits."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== COMBINATIONAL LOGIC ===")

        s, t = self._test_mux2to1(pop, debug)
        scores += s
        total += t

        s, t = self._test_decoder3to8(pop, debug)
        scores += s
        total += t

        s, t = self._test_barrel_shifter(pop, debug)
        scores += s
        total += t

        s, t = self._test_priority_encoder(pop, debug)
        scores += s
        total += t

        return scores, total

    def _test_barrel_shifter(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test barrel shifter (shift by 0-7 positions)."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== BARREL SHIFTER ===")

        try:
            # Test all shift amounts 0-7 with various input patterns
            test_vals = [0b10000001, 0b11110000, 0b00001111, 0b10101010, 0xFF]

            for val in test_vals:
                for shift in range(8):
                    expected_val = (val << shift) & 0xFF  # Left shift
                    shift_bits = [float((shift >> (2 - i)) & 1) for i in range(3)]

                    # Process through 3 layers; every intermediate stays
                    # per-slot ([pop_size]) so population members are
                    # evaluated independently.
                    layer_in = [torch.full((pop_size,), float((val >> (7 - i)) & 1),
                                           device=self.device) for i in range(8)]
                    for layer in range(3):
                        shift_amount = 1 << (2 - layer)  # 4, 2, 1
                        sel = torch.full((pop_size,), shift_bits[layer], device=self.device)
                        layer_out = []

                        for bit in range(8):
                            prefix = f'combinational.barrelshifter.layer{layer}.bit{bit}'

                            # NOT sel
                            w_not = pop[f'{prefix}.not_sel.weight'].view(pop_size)
                            b_not = pop[f'{prefix}.not_sel.bias'].view(pop_size)
                            not_sel = heaviside(sel * w_not + b_not)

                            # Source for shifted value
                            shifted_src = bit + shift_amount
                            if shifted_src < 8:
                                shifted_val = layer_in[shifted_src]
                            else:
                                shifted_val = torch.zeros(pop_size, device=self.device)

                            # AND a: original AND NOT sel
                            w_and_a = pop[f'{prefix}.and_a.weight'].view(pop_size, 2)
                            b_and_a = pop[f'{prefix}.and_a.bias'].view(pop_size)
                            inp_a = torch.stack([layer_in[bit], not_sel], dim=-1)
                            and_a = heaviside((inp_a * w_and_a).sum(-1) + b_and_a)

                            # AND b: shifted AND sel
                            w_and_b = pop[f'{prefix}.and_b.weight'].view(pop_size, 2)
                            b_and_b = pop[f'{prefix}.and_b.bias'].view(pop_size)
                            inp_b = torch.stack([shifted_val, sel], dim=-1)
                            and_b = heaviside((inp_b * w_and_b).sum(-1) + b_and_b)

                            # OR
                            w_or = pop[f'{prefix}.or.weight'].view(pop_size, 2)
                            b_or = pop[f'{prefix}.or.bias'].view(pop_size)
                            inp_or = torch.stack([and_a, and_b], dim=-1)
                            layer_out.append(heaviside((inp_or * w_or).sum(-1) + b_or))

                        layer_in = layer_out

                    # Check result per slot
                    result = torch.zeros(pop_size, device=self.device)
                    for i in range(8):
                        result += layer_in[i] * (1 << (7 - i))
                    scores += (result == expected_val).float()
                    total += 1

            self._record('combinational.barrelshifter', int(scores[0].item()), total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  combinational.barrelshifter: SKIP ({e})")

        return scores, total

    def _test_priority_encoder(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test priority encoder (find highest set bit)."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== PRIORITY ENCODER ===")

        try:
            # Test cases: input -> (valid, index of highest bit)
            test_cases = [
                (0b00000000, 0, 0),  # No bits set, valid=0
                (0b00000001, 1, 7),  # Bit 7 (LSB)
                (0b00000010, 1, 6),
                (0b00000100, 1, 5),
                (0b00001000, 1, 4),
                (0b00010000, 1, 3),
                (0b00100000, 1, 2),
                (0b01000000, 1, 1),
                (0b10000000, 1, 0),  # Bit 0 (MSB)
                (0b10000001, 1, 0),  # Multiple bits, highest wins
                (0b01010101, 1, 1),
                (0b00001111, 1, 4),
                (0b11111111, 1, 0),
            ]

            for val, expected_valid, expected_idx in test_cases:
                val_bits = torch.tensor([float((val >> (7 - i)) & 1) for i in range(8)],
                                       device=self.device, dtype=torch.float32)

                # Valid output: OR of all input bits
                w_valid = pop['combinational.priorityencoder.valid.weight'].view(pop_size, 8)
                b_valid = pop['combinational.priorityencoder.valid.bias'].view(pop_size)
                out_valid = heaviside((val_bits * w_valid).sum(-1) + b_valid)

                scores += (out_valid == float(expected_valid)).float()
                total += 1

                # Index outputs (3 bits)
                if expected_valid == 1:
                    for idx_bit in range(3):
                        try:
                            w_idx = pop[f'combinational.priorityencoder.idx{idx_bit}.weight'].view(pop_size, 8)
                            b_idx = pop[f'combinational.priorityencoder.idx{idx_bit}.bias'].view(pop_size)
                            out_idx = heaviside((val_bits * w_idx).sum(-1) + b_idx)
                            expected_bit = (expected_idx >> (2 - idx_bit)) & 1
                            scores += (out_idx == float(expected_bit)).float()
                            total += 1
                        except KeyError:
                            pass

            self._record('combinational.priorityencoder', int(scores[0].item()), total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  combinational.priorityencoder: SKIP ({e})")

        return scores, total

    def _test_barrel_shifter_nbits(self, pop: Dict, bits: int, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test N-bit barrel shifter (shift by 0 to bits-1 positions)."""
        import math
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0
        num_layers = max(1, math.ceil(math.log2(bits)))
        max_val = (1 << bits) - 1

        if debug:
            print(f"\n=== {bits}-BIT BARREL SHIFTER ===")

        prefix = f'combinational.barrelshifter{bits}'
        try:
            if bits == 16:
                test_vals = [0x8001, 0xFF00, 0x00FF, 0xAAAA, 0xFFFF, 0x1234]
            elif bits == 32:
                test_vals = [0x80000001, 0xFFFF0000, 0x0000FFFF, 0xAAAAAAAA, 0xFFFFFFFF, 0x12345678]
            else:
                test_vals = [0b10000001, 0b11110000, 0b00001111, 0b10101010, max_val]

            num_shifts = min(bits, 8)
            for val in test_vals:
                for shift in range(num_shifts):
                    expected_val = (val << shift) & max_val
                    shift_bits = [float((shift >> (num_layers - 1 - i)) & 1) for i in range(num_layers)]

                    layer_in = [torch.full((pop_size,), float((val >> (bits - 1 - i)) & 1),
                                           device=self.device) for i in range(bits)]
                    for layer in range(num_layers):
                        shift_amount = 1 << (num_layers - 1 - layer)
                        sel = torch.full((pop_size,), shift_bits[layer], device=self.device)
                        layer_out = []

                        for bit in range(bits):
                            bit_prefix = f'{prefix}.layer{layer}.bit{bit}'

                            w_not = pop[f'{bit_prefix}.not_sel.weight'].view(pop_size)
                            b_not = pop[f'{bit_prefix}.not_sel.bias'].view(pop_size)
                            not_sel = heaviside(sel * w_not + b_not)

                            shifted_src = bit + shift_amount
                            if shifted_src < bits:
                                shifted_val = layer_in[shifted_src]
                            else:
                                shifted_val = torch.zeros(pop_size, device=self.device)

                            w_and_a = pop[f'{bit_prefix}.and_a.weight'].view(pop_size, 2)
                            b_and_a = pop[f'{bit_prefix}.and_a.bias'].view(pop_size)
                            inp_a = torch.stack([layer_in[bit], not_sel], dim=-1)
                            and_a = heaviside((inp_a * w_and_a).sum(-1) + b_and_a)

                            w_and_b = pop[f'{bit_prefix}.and_b.weight'].view(pop_size, 2)
                            b_and_b = pop[f'{bit_prefix}.and_b.bias'].view(pop_size)
                            inp_b = torch.stack([shifted_val, sel], dim=-1)
                            and_b = heaviside((inp_b * w_and_b).sum(-1) + b_and_b)

                            w_or = pop[f'{bit_prefix}.or.weight'].view(pop_size, 2)
                            b_or = pop[f'{bit_prefix}.or.bias'].view(pop_size)
                            inp_or = torch.stack([and_a, and_b], dim=-1)
                            layer_out.append(heaviside((inp_or * w_or).sum(-1) + b_or))

                        layer_in = layer_out

                    result = torch.zeros(pop_size, device=self.device, dtype=torch.float64)
                    for i in range(bits):
                        result += layer_in[i].double() * (1 << (bits - 1 - i))
                    scores += (result == float(expected_val)).float()
                    total += 1

            self._record(prefix, int(scores[0].item()), total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  {prefix}: SKIP ({e})")

        return scores, total

    def _test_priority_encoder_nbits(self, pop: Dict, bits: int, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test N-bit priority encoder (find highest set bit).

        The priority encoder is a multi-layer circuit:
        1. any_higher{pos}: OR of bits 0 to pos-1 (all higher-priority positions)
        2. is_highest{0}: bit[0] directly (MSB is always highest if set)
        3. is_highest{pos}: bit[pos] AND NOT(any_higher{pos}) for pos > 0
        4. out{bit}: OR of is_highest{pos} for all pos where (pos >> bit) & 1
        5. valid: OR of all input bits
        """
        import math
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0
        out_bits = max(1, math.ceil(math.log2(bits)))

        if debug:
            print(f"\n=== {bits}-BIT PRIORITY ENCODER ===")

        prefix = f'combinational.priorityencoder{bits}'
        try:
            test_cases = [(0, 0, 0)]
            for i in range(bits):
                test_cases.append((1 << i, 1, bits - 1 - i))
            if bits == 16:
                test_cases.extend([
                    (0x8001, 1, 0), (0x5555, 1, 1), (0x00FF, 1, 8), (0xFFFF, 1, 0)
                ])
            elif bits == 32:
                test_cases.extend([
                    (0x80000001, 1, 0), (0x55555555, 1, 1), (0x0000FFFF, 1, 16), (0xFFFFFFFF, 1, 0)
                ])

            for val, expected_valid, expected_idx in test_cases:
                val_bits = torch.tensor([float((val >> (bits - 1 - i)) & 1) for i in range(bits)],
                                       device=self.device, dtype=torch.float32)

                w_valid = pop[f'{prefix}.valid.weight'].view(pop_size, bits)
                b_valid = pop[f'{prefix}.valid.bias'].view(pop_size)
                out_valid = heaviside((val_bits * w_valid).sum(-1) + b_valid)

                scores += (out_valid == float(expected_valid)).float()
                total += 1

                if expected_valid == 1:
                    any_higher = [None]
                    for pos in range(1, bits):
                        w = pop[f'{prefix}.any_higher{pos}.weight'].view(pop_size, -1)
                        b = pop[f'{prefix}.any_higher{pos}.bias'].view(pop_size)
                        inp = val_bits[:pos]
                        out = heaviside((inp * w[:, :len(inp)]).sum(-1) + b)
                        any_higher.append(out)

                    is_highest = []
                    for pos in range(bits):
                        if pos == 0:
                            is_high = val_bits[0].unsqueeze(0).expand(pop_size)
                        else:
                            w_not = pop[f'{prefix}.is_highest{pos}.not_higher.weight'].view(pop_size)
                            b_not = pop[f'{prefix}.is_highest{pos}.not_higher.bias'].view(pop_size)
                            not_higher = heaviside(any_higher[pos] * w_not + b_not)

                            w_and = pop[f'{prefix}.is_highest{pos}.and.weight'].view(pop_size, -1)
                            b_and = pop[f'{prefix}.is_highest{pos}.and.bias'].view(pop_size)
                            inp = torch.stack([val_bits[pos].expand(pop_size), not_higher], dim=-1)
                            is_high = heaviside((inp * w_and).sum(-1) + b_and)
                        is_highest.append(is_high)

                    for idx_bit in range(out_bits):
                        try:
                            w_idx = pop[f'{prefix}.out{idx_bit}.weight'].view(pop_size, -1)
                            b_idx = pop[f'{prefix}.out{idx_bit}.bias'].view(pop_size)
                            relevant = [is_highest[pos] for pos in range(bits) if (pos >> idx_bit) & 1]
                            if len(relevant) > 0:
                                inp = torch.stack(relevant[:w_idx.shape[1]], dim=-1)
                                out_idx = heaviside((inp * w_idx).sum(-1) + b_idx)
                                expected_bit = (expected_idx >> idx_bit) & 1
                                scores += (out_idx == float(expected_bit)).float()
                                total += 1
                        except KeyError:
                            pass

            self._record(prefix, int(scores[0].item()), total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  {prefix}: SKIP ({e})")

        return scores, total

    # =========================================================================
    # CONTROL FLOW
    # =========================================================================

    def _test_conditional_jump(self, pop: Dict, name: str, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test conditional jump circuit (N-bit address aware)."""
        pop_size = next(iter(pop.values())).shape[0]
        prefix = f'control.{name}'

        # Test cases: [pc_bit, target_bit, flag] -> out = flag ? target : pc
        inputs = torch.tensor([
            [0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1],
            [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1],
        ], device=self.device, dtype=torch.float32)
        expected = torch.tensor([0, 0, 0, 1, 1, 0, 1, 1], device=self.device, dtype=torch.float32)

        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        pc = inputs[:, 0].unsqueeze(1).expand(-1, pop_size)      # [8, pop]
        target = inputs[:, 1].unsqueeze(1).expand(-1, pop_size)
        flag = inputs[:, 2].unsqueeze(1).expand(-1, pop_size)

        for bit in range(self.addr_bits):
            bit_prefix = f'{prefix}.bit{bit}'
            try:
                # NOT sel
                w_not = pop[f'{bit_prefix}.not_sel.weight'].view(pop_size)
                b_not = pop[f'{bit_prefix}.not_sel.bias'].view(pop_size)
                not_sel = heaviside(flag * w_not + b_not)

                # AND a (pc AND NOT sel)
                w_and_a = pop[f'{bit_prefix}.and_a.weight'].view(pop_size, 2)
                b_and_a = pop[f'{bit_prefix}.and_a.bias'].view(pop_size)
                inp_a = torch.stack([pc, not_sel], dim=-1)
                and_a = heaviside((inp_a * w_and_a).sum(-1) + b_and_a)

                # AND b (target AND sel)
                w_and_b = pop[f'{bit_prefix}.and_b.weight'].view(pop_size, 2)
                b_and_b = pop[f'{bit_prefix}.and_b.bias'].view(pop_size)
                inp_b = torch.stack([target, flag], dim=-1)
                and_b = heaviside((inp_b * w_and_b).sum(-1) + b_and_b)

                # OR
                w_or = pop[f'{bit_prefix}.or.weight'].view(pop_size, 2)
                b_or = pop[f'{bit_prefix}.or.bias'].view(pop_size)
                ab = torch.stack([and_a, and_b], dim=-1)  # [8, pop_size, 2]
                out = heaviside((ab * w_or).sum(-1) + b_or)  # [8, pop_size]

                correct = (out == expected.unsqueeze(1)).float().sum(0)  # [pop_size]
                scores += correct
                total += 8

            except KeyError:
                pass

        if total > 0:
            self._record(prefix, int(scores[0].item()), total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")

        return scores, total

    def _test_control_flow(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test control flow circuits."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== CONTROL FLOW ===")

        jumps = ['jz', 'jnz', 'jc', 'jnc', 'jn', 'jp', 'jv', 'jnv', 'conditionaljump']
        for name in jumps:
            s, t = self._test_conditional_jump(pop, name, debug)
            scores += s
            total += t

        # Stack operations
        s, t = self._test_stack_ops(pop, debug)
        scores += s
        total += t

        return scores, total

    def _test_stack_ops(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test PUSH/POP/RET stack operation circuits (N-bit address aware)."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0
        addr_bits = self.addr_bits
        addr_mask = (1 << addr_bits) - 1

        if debug:
            print(f"\n=== STACK OPERATIONS ({addr_bits}-bit SP) ===")

        # Test PUSH SP decrement (addr_bits wide, borrow chain)
        try:
            # Generate test values appropriate for addr_bits
            sp_tests = [0, 1, addr_mask // 2, addr_mask]
            if addr_bits >= 8:
                sp_tests.append(0x100 & addr_mask)
            if addr_bits >= 12:
                sp_tests.append(0x1234 & addr_mask)
            op_scores = torch.zeros(pop_size, device=self.device)
            op_total = 0

            for sp_val in sp_tests:
                expected_val = (sp_val - 1) & addr_mask
                sp_bits = [float((sp_val >> (addr_bits - 1 - i)) & 1) for i in range(addr_bits)]

                borrow = torch.ones(pop_size, device=self.device)
                out_bits = []
                for bit in range(addr_bits - 1, -1, -1):  # LSB to MSB
                    prefix = f'control.push.sp_dec.bit{bit}'

                    w_or = pop[f'{prefix}.xor.layer1.or.weight'].view(pop_size, 2)
                    b_or = pop[f'{prefix}.xor.layer1.or.bias'].view(pop_size)
                    w_nand = pop[f'{prefix}.xor.layer1.nand.weight'].view(pop_size, 2)
                    b_nand = pop[f'{prefix}.xor.layer1.nand.bias'].view(pop_size)
                    w2 = pop[f'{prefix}.xor.layer2.weight'].view(pop_size, 2)
                    b2 = pop[f'{prefix}.xor.layer2.bias'].view(pop_size)

                    sp_bit = torch.full((pop_size,), sp_bits[bit], device=self.device)
                    inp = torch.stack([sp_bit, borrow], dim=-1)
                    h_or = heaviside((inp * w_or).sum(-1) + b_or)
                    h_nand = heaviside((inp * w_nand).sum(-1) + b_nand)
                    hidden = torch.stack([h_or, h_nand], dim=-1)
                    diff_bit = heaviside((hidden * w2).sum(-1) + b2)
                    out_bits.insert(0, diff_bit)

                    # Borrow: NOT(sp) AND borrow_in
                    not_sp = torch.full((pop_size,), 1.0 - sp_bits[bit], device=self.device)
                    w_borrow = pop[f'{prefix}.borrow.weight'].view(pop_size, 2)
                    b_borrow = pop[f'{prefix}.borrow.bias'].view(pop_size)
                    borrow_inp = torch.stack([not_sp, borrow], dim=-1)
                    borrow = heaviside((borrow_inp * w_borrow).sum(-1) + b_borrow)

                out = torch.stack(out_bits, dim=-1)
                expected = torch.tensor([((expected_val >> (addr_bits - 1 - i)) & 1) for i in range(addr_bits)],
                                       device=self.device, dtype=torch.float32)
                correct = (out == expected.unsqueeze(0)).float().sum(1)
                op_scores += correct
                op_total += addr_bits

            scores += op_scores
            total += op_total
            self._record('control.push.sp_dec', int(op_scores[0].item()), op_total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  control.push.sp_dec: SKIP ({e})")

        # Test POP SP increment (addr_bits wide, carry chain)
        try:
            op_scores = torch.zeros(pop_size, device=self.device)
            op_total = 0

            for sp_val in sp_tests:
                expected_val = (sp_val + 1) & addr_mask
                sp_bits = [float((sp_val >> (addr_bits - 1 - i)) & 1) for i in range(addr_bits)]

                carry = torch.ones(pop_size, device=self.device)
                out_bits = []
                for bit in range(addr_bits - 1, -1, -1):  # LSB to MSB
                    prefix = f'control.pop.sp_inc.bit{bit}'

                    w_or = pop[f'{prefix}.xor.layer1.or.weight'].view(pop_size, 2)
                    b_or = pop[f'{prefix}.xor.layer1.or.bias'].view(pop_size)
                    w_nand = pop[f'{prefix}.xor.layer1.nand.weight'].view(pop_size, 2)
                    b_nand = pop[f'{prefix}.xor.layer1.nand.bias'].view(pop_size)
                    w2 = pop[f'{prefix}.xor.layer2.weight'].view(pop_size, 2)
                    b2 = pop[f'{prefix}.xor.layer2.bias'].view(pop_size)

                    sp_bit = torch.full((pop_size,), sp_bits[bit], device=self.device)
                    inp = torch.stack([sp_bit, carry], dim=-1)
                    h_or = heaviside((inp * w_or).sum(-1) + b_or)
                    h_nand = heaviside((inp * w_nand).sum(-1) + b_nand)
                    hidden = torch.stack([h_or, h_nand], dim=-1)
                    sum_bit = heaviside((hidden * w2).sum(-1) + b2)
                    out_bits.insert(0, sum_bit)

                    # Carry: sp AND carry_in
                    w_carry = pop[f'{prefix}.carry.weight'].view(pop_size, 2)
                    b_carry = pop[f'{prefix}.carry.bias'].view(pop_size)
                    carry = heaviside((inp * w_carry).sum(-1) + b_carry)

                out = torch.stack(out_bits, dim=-1)
                expected = torch.tensor([((expected_val >> (addr_bits - 1 - i)) & 1) for i in range(addr_bits)],
                                       device=self.device, dtype=torch.float32)
                correct = (out == expected.unsqueeze(0)).float().sum(1)
                op_scores += correct
                op_total += addr_bits

            scores += op_scores
            total += op_total
            self._record('control.pop.sp_inc', int(op_scores[0].item()), op_total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  control.pop.sp_inc: SKIP ({e})")

        # Test RET address buffer (addr_bits identity gates)
        try:
            op_scores = torch.zeros(pop_size, device=self.device)
            op_total = 0

            ret_tests = [0, addr_mask, addr_mask // 2, 1]
            if addr_bits >= 12:
                ret_tests.append(0x1234 & addr_mask)
            for addr_val in ret_tests:
                ret_bits_tensor = torch.tensor([float((addr_val >> (addr_bits - 1 - i)) & 1) for i in range(addr_bits)],
                                        device=self.device, dtype=torch.float32)

                out_bits = []
                for bit in range(addr_bits):
                    w = pop[f'control.ret.addr.bit{bit}.weight'].view(pop_size)
                    b = pop[f'control.ret.addr.bit{bit}.bias'].view(pop_size)
                    out = heaviside(ret_bits_tensor[bit] * w + b)
                    out_bits.append(out)

                out = torch.stack(out_bits, dim=-1)
                correct = (out == ret_bits_tensor.unsqueeze(0)).float().sum(1)
                op_scores += correct
                op_total += addr_bits

            scores += op_scores
            total += op_total
            self._record('control.ret.addr', int(op_scores[0].item()), op_total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  control.ret.addr: SKIP ({e})")

        return scores, total

    # =========================================================================
    # ALU
    # =========================================================================

    def _test_alu_ops(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test ALU operations (8-bit bitwise)."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== ALU OPERATIONS ===")

        # Test ALU AND/OR/NOT on 8-bit values
        # Each ALU op has weight [16] or [8] and bias [8]
        # Structured as 8 parallel 2-input (or 1-input for NOT) gates

        test_vals = [(0, 0), (255, 255), (0xAA, 0x55), (0x0F, 0xF0)]

        # AND: weight [16] = 8 * [2], bias [8]
        try:
            w = pop['alu.alu8bit.and.weight'].view(pop_size, 8, 2)  # [pop, 8, 2]
            b = pop['alu.alu8bit.and.bias'].view(pop_size, 8)  # [pop, 8]
            op_scores = torch.zeros(pop_size, device=self.device)
            op_total = 0

            for a_val, b_val in test_vals:
                a_bits = torch.tensor([((a_val >> (7 - i)) & 1) for i in range(8)],
                                     device=self.device, dtype=torch.float32)
                b_bits = torch.tensor([((b_val >> (7 - i)) & 1) for i in range(8)],
                                     device=self.device, dtype=torch.float32)
                # [8, 2]
                inputs = torch.stack([a_bits, b_bits], dim=-1)
                # [pop, 8]
                out = heaviside((inputs * w).sum(-1) + b)
                expected = torch.tensor([((a_val & b_val) >> (7 - i)) & 1 for i in range(8)],
                                       device=self.device, dtype=torch.float32)
                correct = (out == expected.unsqueeze(0)).float().sum(1)  # [pop]
                op_scores += correct
                op_total += 8

            scores += op_scores
            total += op_total
            self._record('alu.alu8bit.and', int(op_scores[0].item()), op_total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError):
            pass

        # OR
        try:
            w = pop['alu.alu8bit.or.weight'].view(pop_size, 8, 2)
            b = pop['alu.alu8bit.or.bias'].view(pop_size, 8)
            op_scores = torch.zeros(pop_size, device=self.device)
            op_total = 0

            for a_val, b_val in test_vals:
                a_bits = torch.tensor([((a_val >> (7 - i)) & 1) for i in range(8)],
                                     device=self.device, dtype=torch.float32)
                b_bits = torch.tensor([((b_val >> (7 - i)) & 1) for i in range(8)],
                                     device=self.device, dtype=torch.float32)
                inputs = torch.stack([a_bits, b_bits], dim=-1)
                out = heaviside((inputs * w).sum(-1) + b)
                expected = torch.tensor([((a_val | b_val) >> (7 - i)) & 1 for i in range(8)],
                                       device=self.device, dtype=torch.float32)
                correct = (out == expected.unsqueeze(0)).float().sum(1)
                op_scores += correct
                op_total += 8

            scores += op_scores
            total += op_total
            self._record('alu.alu8bit.or', int(op_scores[0].item()), op_total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError):
            pass

        # NOT
        try:
            w = pop['alu.alu8bit.not.weight'].view(pop_size, 8)
            b = pop['alu.alu8bit.not.bias'].view(pop_size, 8)
            op_scores = torch.zeros(pop_size, device=self.device)
            op_total = 0

            for a_val, _ in test_vals:
                a_bits = torch.tensor([((a_val >> (7 - i)) & 1) for i in range(8)],
                                     device=self.device, dtype=torch.float32)
                out = heaviside(a_bits * w + b)
                expected = torch.tensor([(((~a_val) & 0xFF) >> (7 - i)) & 1 for i in range(8)],
                                       device=self.device, dtype=torch.float32)
                correct = (out == expected.unsqueeze(0)).float().sum(1)
                op_scores += correct
                op_total += 8

            scores += op_scores
            total += op_total
            self._record('alu.alu8bit.not', int(op_scores[0].item()), op_total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError):
            pass

        # SHL (shift left)
        try:
            op_scores = torch.zeros(pop_size, device=self.device)
            op_total = 0

            for a_val, _ in test_vals:
                expected_val = (a_val << 1) & 0xFF
                a_bits = torch.tensor([((a_val >> (7 - i)) & 1) for i in range(8)],
                                     device=self.device, dtype=torch.float32)
                out_bits = []
                for bit in range(8):
                    w = pop[f'alu.alu8bit.shl.bit{bit}.weight'].view(pop_size)
                    b = pop[f'alu.alu8bit.shl.bit{bit}.bias'].view(pop_size)
                    if bit < 7:
                        inp = a_bits[bit + 1].unsqueeze(0).expand(pop_size)
                    else:
                        inp = torch.zeros(pop_size, device=self.device)
                    out = heaviside(inp * w + b)
                    out_bits.append(out)
                out = torch.stack(out_bits, dim=-1)  # [pop, 8]
                expected = torch.tensor([((expected_val >> (7 - i)) & 1) for i in range(8)],
                                       device=self.device, dtype=torch.float32)
                correct = (out == expected.unsqueeze(0)).float().sum(1)
                op_scores += correct
                op_total += 8

            scores += op_scores
            total += op_total
            self._record('alu.alu8bit.shl', int(op_scores[0].item()), op_total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  alu.alu8bit.shl: SKIP ({e})")

        # SHR (shift right)
        try:
            op_scores = torch.zeros(pop_size, device=self.device)
            op_total = 0

            for a_val, _ in test_vals:
                expected_val = (a_val >> 1) & 0xFF
                a_bits = torch.tensor([((a_val >> (7 - i)) & 1) for i in range(8)],
                                     device=self.device, dtype=torch.float32)
                out_bits = []
                for bit in range(8):
                    w = pop[f'alu.alu8bit.shr.bit{bit}.weight'].view(pop_size)
                    b = pop[f'alu.alu8bit.shr.bit{bit}.bias'].view(pop_size)
                    if bit > 0:
                        inp = a_bits[bit - 1].unsqueeze(0).expand(pop_size)
                    else:
                        inp = torch.zeros(pop_size, device=self.device)
                    out = heaviside(inp * w + b)
                    out_bits.append(out)
                out = torch.stack(out_bits, dim=-1)  # [pop, 8]
                expected = torch.tensor([((expected_val >> (7 - i)) & 1) for i in range(8)],
                                       device=self.device, dtype=torch.float32)
                correct = (out == expected.unsqueeze(0)).float().sum(1)
                op_scores += correct
                op_total += 8

            scores += op_scores
            total += op_total
            self._record('alu.alu8bit.shr', int(op_scores[0].item()), op_total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  alu.alu8bit.shr: SKIP ({e})")

        # MUL (partial products only - just verify AND gates work)
        try:
            op_scores = torch.zeros(pop_size, device=self.device)
            op_total = 0

            mul_tests = [(3, 4), (7, 8), (15, 17), (0, 255)]
            for a_val, b_val in mul_tests:
                a_bits = torch.tensor([((a_val >> (7 - i)) & 1) for i in range(8)],
                                     device=self.device, dtype=torch.float32)
                b_bits = torch.tensor([((b_val >> (7 - i)) & 1) for i in range(8)],
                                     device=self.device, dtype=torch.float32)

                # Test partial product AND gates
                for i in range(8):
                    for j in range(8):
                        w = pop[f'alu.alu8bit.mul.pp.a{i}b{j}.weight'].view(pop_size, 2)
                        b = pop[f'alu.alu8bit.mul.pp.a{i}b{j}.bias'].view(pop_size)
                        inp = torch.tensor([a_bits[i].item(), b_bits[j].item()], device=self.device)
                        out = heaviside((inp * w).sum(-1) + b)
                        expected = float(int(a_bits[i].item()) & int(b_bits[j].item()))
                        correct = (out == expected).float()
                        op_scores += correct
                        op_total += 1

            scores += op_scores
            total += op_total
            self._record('alu.alu8bit.mul', int(op_scores[0].item()), op_total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  alu.alu8bit.mul: SKIP ({e})")

        # DIV: drive each stage's bit-cascade GE comparator along the real
        # restoring-division remainder trace for each operand pair.
        try:
            op_scores = torch.zeros(pop_size, device=self.device)
            op_total = 0

            div_tests = [(100, 10), (255, 17), (50, 7), (128, 16), (255, 255),
                         (254, 255), (7, 130), (200, 3), (9, 3), (1, 1)]
            # Remainder value and expected GE at each stage, per pair.
            stage_rems = [[] for _ in range(8)]
            stage_ges = [[] for _ in range(8)]
            for a_val, b_val in div_tests:
                a_bits_int = [(a_val >> (7 - i)) & 1 for i in range(8)]
                remainder = 0
                for stage in range(8):
                    remainder = ((remainder << 1) | a_bits_int[stage]) & 0xFF
                    stage_rems[stage].append(remainder)
                    ge = 1.0 if remainder >= b_val else 0.0
                    stage_ges[stage].append(ge)
                    if ge:
                        remainder -= b_val

            div_vals = torch.tensor([b for _, b in div_tests], device=self.device, dtype=torch.long)
            div_bits = torch.stack([((div_vals >> (7 - i)) & 1).float() for i in range(8)], dim=1)
            for stage in range(8):
                rem_vals = torch.tensor(stage_rems[stage], device=self.device, dtype=torch.long)
                rem_bits = torch.stack([((rem_vals >> (7 - i)) & 1).float() for i in range(8)], dim=1)
                outs = self._eval_bit_cascade_compare(
                    pop,
                    f'alu.alu8bit.div.stage{stage}.cmp_bc',
                    f'alu.alu8bit.div.stage{stage}.cmp_bc.gt',
                    f'alu.alu8bit.div.stage{stage}.cmp_bc.lt',
                    f'alu.alu8bit.div.stage{stage}.cmp',
                    f'alu.alu8bit.div.stage{stage}.cmp_bc.le',
                    f'alu.alu8bit.div.stage{stage}.cmp_bc.eq',
                    8, rem_bits, div_bits,
                )
                expected = torch.tensor(stage_ges[stage], device=self.device)
                correct = (outs['ge'] == expected.unsqueeze(1)).float().sum(0)  # [pop]
                op_scores += correct
                op_total += len(div_tests)

            scores += op_scores
            total += op_total
            self._record('alu.alu8bit.div', int(op_scores[0].item()), op_total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  alu.alu8bit.div: SKIP ({e})")

        # INC (increment by 1)
        try:
            op_scores = torch.zeros(pop_size, device=self.device)
            op_total = 0

            inc_tests = [0, 1, 127, 128, 254, 255]
            for a_val in inc_tests:
                expected_val = (a_val + 1) & 0xFF
                a_bits = torch.tensor([((a_val >> (7 - i)) & 1) for i in range(8)],
                                     device=self.device, dtype=torch.float32)

                # INC uses half-adder chain with initial carry = 1
                carry = torch.ones(pop_size, device=self.device)
                out_bits = []
                for bit in range(7, -1, -1):  # LSB to MSB
                    # XOR for sum
                    w_or = pop[f'alu.alu8bit.inc.bit{bit}.xor.layer1.or.weight'].view(pop_size, 2)
                    b_or = pop[f'alu.alu8bit.inc.bit{bit}.xor.layer1.or.bias'].view(pop_size)
                    w_nand = pop[f'alu.alu8bit.inc.bit{bit}.xor.layer1.nand.weight'].view(pop_size, 2)
                    b_nand = pop[f'alu.alu8bit.inc.bit{bit}.xor.layer1.nand.bias'].view(pop_size)
                    w2 = pop[f'alu.alu8bit.inc.bit{bit}.xor.layer2.weight'].view(pop_size, 2)
                    b2 = pop[f'alu.alu8bit.inc.bit{bit}.xor.layer2.bias'].view(pop_size)

                    inp = torch.stack([a_bits[bit].expand(pop_size), carry], dim=-1)
                    h_or = heaviside((inp * w_or).sum(-1) + b_or)
                    h_nand = heaviside((inp * w_nand).sum(-1) + b_nand)
                    hidden = torch.stack([h_or, h_nand], dim=-1)
                    sum_bit = heaviside((hidden * w2).sum(-1) + b2)
                    out_bits.insert(0, sum_bit)

                    # AND for carry
                    w_carry = pop[f'alu.alu8bit.inc.bit{bit}.carry.weight'].view(pop_size, 2)
                    b_carry = pop[f'alu.alu8bit.inc.bit{bit}.carry.bias'].view(pop_size)
                    carry = heaviside((inp * w_carry).sum(-1) + b_carry)

                out = torch.stack(out_bits, dim=-1)
                expected = torch.tensor([((expected_val >> (7 - i)) & 1) for i in range(8)],
                                       device=self.device, dtype=torch.float32)
                correct = (out == expected.unsqueeze(0)).float().sum(1)
                op_scores += correct
                op_total += 8

            scores += op_scores
            total += op_total
            self._record('alu.alu8bit.inc', int(op_scores[0].item()), op_total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  alu.alu8bit.inc: SKIP ({e})")

        # DEC (decrement by 1)
        try:
            op_scores = torch.zeros(pop_size, device=self.device)
            op_total = 0

            dec_tests = [0, 1, 127, 128, 254, 255]
            for a_val in dec_tests:
                expected_val = (a_val - 1) & 0xFF
                a_bits = torch.tensor([((a_val >> (7 - i)) & 1) for i in range(8)],
                                     device=self.device, dtype=torch.float32)

                # DEC uses borrow chain
                borrow = torch.ones(pop_size, device=self.device)
                out_bits = []
                for bit in range(7, -1, -1):
                    w_or = pop[f'alu.alu8bit.dec.bit{bit}.xor.layer1.or.weight'].view(pop_size, 2)
                    b_or = pop[f'alu.alu8bit.dec.bit{bit}.xor.layer1.or.bias'].view(pop_size)
                    w_nand = pop[f'alu.alu8bit.dec.bit{bit}.xor.layer1.nand.weight'].view(pop_size, 2)
                    b_nand = pop[f'alu.alu8bit.dec.bit{bit}.xor.layer1.nand.bias'].view(pop_size)
                    w2 = pop[f'alu.alu8bit.dec.bit{bit}.xor.layer2.weight'].view(pop_size, 2)
                    b2 = pop[f'alu.alu8bit.dec.bit{bit}.xor.layer2.bias'].view(pop_size)

                    inp = torch.stack([a_bits[bit].expand(pop_size), borrow], dim=-1)
                    h_or = heaviside((inp * w_or).sum(-1) + b_or)
                    h_nand = heaviside((inp * w_nand).sum(-1) + b_nand)
                    hidden = torch.stack([h_or, h_nand], dim=-1)
                    diff_bit = heaviside((hidden * w2).sum(-1) + b2)
                    out_bits.insert(0, diff_bit)

                    # Borrow logic: borrow_out = NOT(a) AND borrow_in
                    w_not = pop[f'alu.alu8bit.dec.bit{bit}.not_a.weight'].view(pop_size)
                    b_not = pop[f'alu.alu8bit.dec.bit{bit}.not_a.bias'].view(pop_size)
                    not_a = heaviside(a_bits[bit] * w_not + b_not)

                    w_borrow = pop[f'alu.alu8bit.dec.bit{bit}.borrow.weight'].view(pop_size, 2)
                    b_borrow = pop[f'alu.alu8bit.dec.bit{bit}.borrow.bias'].view(pop_size)
                    borrow_inp = torch.stack([not_a, borrow], dim=-1)
                    borrow = heaviside((borrow_inp * w_borrow).sum(-1) + b_borrow)

                out = torch.stack(out_bits, dim=-1)
                expected = torch.tensor([((expected_val >> (7 - i)) & 1) for i in range(8)],
                                       device=self.device, dtype=torch.float32)
                correct = (out == expected.unsqueeze(0)).float().sum(1)
                op_scores += correct
                op_total += 8

            scores += op_scores
            total += op_total
            self._record('alu.alu8bit.dec', int(op_scores[0].item()), op_total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  alu.alu8bit.dec: SKIP ({e})")

        # NEG (two's complement: NOT + 1)
        try:
            op_scores = torch.zeros(pop_size, device=self.device)
            op_total = 0

            neg_tests = [0, 1, 127, 128, 255]
            for a_val in neg_tests:
                expected_val = (-a_val) & 0xFF
                a_bits = torch.tensor([((a_val >> (7 - i)) & 1) for i in range(8)],
                                     device=self.device, dtype=torch.float32)

                # First NOT each bit
                not_bits = []
                for bit in range(8):
                    w = pop[f'alu.alu8bit.neg.not.bit{bit}.weight'].view(pop_size)
                    b = pop[f'alu.alu8bit.neg.not.bit{bit}.bias'].view(pop_size)
                    not_bit = heaviside(a_bits[bit] * w + b)
                    not_bits.append(not_bit)

                # Then INC
                carry = torch.ones(pop_size, device=self.device)
                out_bits = []
                for bit in range(7, -1, -1):
                    w_or = pop[f'alu.alu8bit.neg.inc.bit{bit}.xor.layer1.or.weight'].view(pop_size, 2)
                    b_or = pop[f'alu.alu8bit.neg.inc.bit{bit}.xor.layer1.or.bias'].view(pop_size)
                    w_nand = pop[f'alu.alu8bit.neg.inc.bit{bit}.xor.layer1.nand.weight'].view(pop_size, 2)
                    b_nand = pop[f'alu.alu8bit.neg.inc.bit{bit}.xor.layer1.nand.bias'].view(pop_size)
                    w2 = pop[f'alu.alu8bit.neg.inc.bit{bit}.xor.layer2.weight'].view(pop_size, 2)
                    b2 = pop[f'alu.alu8bit.neg.inc.bit{bit}.xor.layer2.bias'].view(pop_size)

                    inp = torch.stack([not_bits[bit], carry], dim=-1)
                    h_or = heaviside((inp * w_or).sum(-1) + b_or)
                    h_nand = heaviside((inp * w_nand).sum(-1) + b_nand)
                    hidden = torch.stack([h_or, h_nand], dim=-1)
                    sum_bit = heaviside((hidden * w2).sum(-1) + b2)
                    out_bits.insert(0, sum_bit)

                    w_carry = pop[f'alu.alu8bit.neg.inc.bit{bit}.carry.weight'].view(pop_size, 2)
                    b_carry = pop[f'alu.alu8bit.neg.inc.bit{bit}.carry.bias'].view(pop_size)
                    carry = heaviside((inp * w_carry).sum(-1) + b_carry)

                out = torch.stack(out_bits, dim=-1)
                expected = torch.tensor([((expected_val >> (7 - i)) & 1) for i in range(8)],
                                       device=self.device, dtype=torch.float32)
                correct = (out == expected.unsqueeze(0)).float().sum(1)
                op_scores += correct
                op_total += 8

            scores += op_scores
            total += op_total
            self._record('alu.alu8bit.neg', int(op_scores[0].item()), op_total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  alu.alu8bit.neg: SKIP ({e})")

        # ROL (rotate left - MSB wraps to LSB)
        try:
            op_scores = torch.zeros(pop_size, device=self.device)
            op_total = 0

            rol_tests = [0b10000000, 0b00000001, 0b10101010, 0b01010101, 0xFF, 0x00]
            for a_val in rol_tests:
                expected_val = ((a_val << 1) | (a_val >> 7)) & 0xFF
                a_bits = torch.tensor([((a_val >> (7 - i)) & 1) for i in range(8)],
                                     device=self.device, dtype=torch.float32)

                out_bits = []
                for bit in range(8):
                    w = pop[f'alu.alu8bit.rol.bit{bit}.weight'].view(pop_size)
                    b = pop[f'alu.alu8bit.rol.bit{bit}.bias'].view(pop_size)
                    # ROL: bit[i] gets bit[i+1], bit[7] gets bit[0]
                    src_bit = (bit + 1) % 8
                    out = heaviside(a_bits[src_bit] * w + b)
                    out_bits.append(out)

                out = torch.stack(out_bits, dim=-1)
                expected = torch.tensor([((expected_val >> (7 - i)) & 1) for i in range(8)],
                                       device=self.device, dtype=torch.float32)
                correct = (out == expected.unsqueeze(0)).float().sum(1)
                op_scores += correct
                op_total += 8

            scores += op_scores
            total += op_total
            self._record('alu.alu8bit.rol', int(op_scores[0].item()), op_total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  alu.alu8bit.rol: SKIP ({e})")

        # ROR (rotate right - LSB wraps to MSB)
        try:
            op_scores = torch.zeros(pop_size, device=self.device)
            op_total = 0

            ror_tests = [0b10000000, 0b00000001, 0b10101010, 0b01010101, 0xFF, 0x00]
            for a_val in ror_tests:
                expected_val = ((a_val >> 1) | (a_val << 7)) & 0xFF
                a_bits = torch.tensor([((a_val >> (7 - i)) & 1) for i in range(8)],
                                     device=self.device, dtype=torch.float32)

                out_bits = []
                for bit in range(8):
                    w = pop[f'alu.alu8bit.ror.bit{bit}.weight'].view(pop_size)
                    b = pop[f'alu.alu8bit.ror.bit{bit}.bias'].view(pop_size)
                    # ROR: bit[i] gets bit[i-1], bit[0] gets bit[7]
                    src_bit = (bit - 1) % 8
                    out = heaviside(a_bits[src_bit] * w + b)
                    out_bits.append(out)

                out = torch.stack(out_bits, dim=-1)
                expected = torch.tensor([((expected_val >> (7 - i)) & 1) for i in range(8)],
                                       device=self.device, dtype=torch.float32)
                correct = (out == expected.unsqueeze(0)).float().sum(1)
                op_scores += correct
                op_total += 8

            scores += op_scores
            total += op_total
            self._record('alu.alu8bit.ror', int(op_scores[0].item()), op_total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  alu.alu8bit.ror: SKIP ({e})")

        return scores, total

    # =========================================================================
    # MANIFEST
    # =========================================================================

    def _test_manifest(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Verify manifest values."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== MANIFEST ===")

        fixed_expected = {
            'manifest.alu_operations': 16.0,
            'manifest.flags': 4.0,
            'manifest.instruction_width': 16.0,
            'manifest.register_width': 8.0,
            'manifest.registers': 4.0,
            'manifest.version': 4.0,
        }

        for name, exp_val in fixed_expected.items():
            try:
                val = pop[name][0, 0].item()
                if val == exp_val:
                    scores += 1
                    self._record(name, 1, 1, [])
                else:
                    self._record(name, 0, 1, [(exp_val, val)])
                total += 1
                if debug:
                    r = self.results[-1]
                    print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
            except KeyError:
                pass

        variable_checks = ['manifest.memory_bytes', 'manifest.pc_width', 'manifest.turing_complete']
        for name in variable_checks:
            try:
                val = pop[name][0, 0].item()
                valid = val >= 0
                if valid:
                    scores += 1
                    self._record(name, 1, 1, [])
                else:
                    self._record(name, 0, 1, [('>=0', val)])
                total += 1
                if debug:
                    r = self.results[-1]
                    print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'} (value={val})")
            except KeyError:
                pass

        return scores, total

    # =========================================================================
    # MEMORY
    # =========================================================================

    def _test_memory(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test memory circuits (shape validation)."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== MEMORY ===")

        try:
            mem_bytes = int(pop['manifest.memory_bytes'][0].item())
            addr_bits = int(pop['manifest.pc_width'][0].item())
        except KeyError:
            mem_bytes = 65536
            addr_bits = 16

        if mem_bytes == 0:
            if debug:
                print("  No memory (pure ALU mode)")
            return scores, 0

        expected_shapes = {
            'memory.addr_decode.weight': (mem_bytes, addr_bits),
            'memory.addr_decode.bias': (mem_bytes,),
            'memory.read.and.weight': (8, mem_bytes, 2),
            'memory.read.and.bias': (8, mem_bytes),
            'memory.read.or.weight': (8, mem_bytes),
            'memory.read.or.bias': (8,),
            'memory.write.sel.weight': (mem_bytes, 2),
            'memory.write.sel.bias': (mem_bytes,),
            'memory.write.nsel.weight': (mem_bytes, 1),
            'memory.write.nsel.bias': (mem_bytes,),
            'memory.write.and_old.weight': (mem_bytes, 8, 2),
            'memory.write.and_old.bias': (mem_bytes, 8),
            'memory.write.and_new.weight': (mem_bytes, 8, 2),
            'memory.write.and_new.bias': (mem_bytes, 8),
            'memory.write.or.weight': (mem_bytes, 8, 2),
            'memory.write.or.bias': (mem_bytes, 8),
        }

        for name, expected_shape in expected_shapes.items():
            try:
                tensor = pop[name]
                actual_shape = tuple(tensor.shape[1:])  # Skip pop_size dimension
                if actual_shape == expected_shape:
                    scores += 1
                    self._record(name, 1, 1, [])
                else:
                    self._record(name, 0, 1, [(expected_shape, actual_shape)])
                total += 1

                if debug:
                    r = self.results[-1]
                    print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
            except KeyError:
                pass

        return scores, total

    # =========================================================================
    # FLOAT TESTS
    #
    # unpack/pack buffers and the classify subcircuit are functionally tested
    # (inputs driven through the gates, outputs compared to IEEE 754
    # semantics). The composed add/mul/div/cmp pipelines are self-contained:
    # their full wiring ships as .inputs metadata, so each is reconstructed
    # with NetlistEvaluator and evaluated end to end against exact integer
    # oracles (round-to-nearest-even, bit-exact to IEEE hardware). The
    # remaining per-stage shape checks below just confirm the gate inventory.
    # =========================================================================

    def _test_float_unpack_pack(self, pop: Dict, family: str, word_bits: int,
                                debug: bool) -> Tuple[torch.Tensor, int]:
        """Functionally test the unpack/pack identity buffers: every bit gate
        must reproduce its binary input (0 -> 0, 1 -> 1)."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print(f"\n=== {family.upper()} UNPACK/PACK (functional) ===")

        inputs = torch.tensor([[0.0], [1.0]], device=self.device)
        expected = torch.tensor([0.0, 1.0], device=self.device)
        for stage in ("unpack", "pack"):
            try:
                ok = torch.zeros(pop_size, device=self.device)
                t = 0
                for i in range(word_bits):
                    w = pop[f'{family}.{stage}.bit{i}.weight'].view(pop_size, 1)
                    b = pop[f'{family}.{stage}.bit{i}.bias'].view(pop_size)
                    out = heaviside(inputs @ w.T + b)  # [2, pop]
                    ok += (out == expected.unsqueeze(1)).float().sum(0)
                    t += 2
            except KeyError:
                continue
            scores += ok
            total += t
            self._record(f'{family}.{stage}', int(ok[0].item()), t, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")

        return scores, total

    def _test_float_classify(self, pop: Dict, family: str, exp_bits: int,
                             frac_bits: int, debug: bool) -> Tuple[torch.Tensor, int]:
        """Functionally test the classify subcircuit: the exponent/fraction
        field predicates and the is_zero / is_subnormal / is_inf / is_nan AND
        gates, against IEEE 754 categories over edge-case encodings."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print(f"\n=== {family.upper()} CLASSIFY (functional) ===")

        exp_max_val = (1 << exp_bits) - 1
        frac_max_val = (1 << frac_bits) - 1
        # (exponent, fraction) pairs covering every IEEE category at both extremes.
        cases = [
            (0, 0),                            # zero
            (0, 1),                            # smallest subnormal
            (0, frac_max_val),                 # largest subnormal
            (1, 0),                            # smallest normal
            (exp_max_val - 1, frac_max_val),   # largest normal
            (exp_max_val // 2, 0),             # mid-range normal
            (exp_max_val, 0),                  # infinity
            (exp_max_val, 1),                  # NaN, minimal payload
            (exp_max_val, frac_max_val),       # NaN, full payload
        ]
        num = len(cases)
        exp_vals = torch.tensor([c[0] for c in cases], device=self.device, dtype=torch.long)
        frac_vals = torch.tensor([c[1] for c in cases], device=self.device, dtype=torch.long)
        exp_in = torch.stack([((exp_vals >> (exp_bits - 1 - i)) & 1).float()
                              for i in range(exp_bits)], dim=1)
        frac_in = torch.stack([((frac_vals >> (frac_bits - 1 - i)) & 1).float()
                               for i in range(frac_bits)], dim=1)

        try:
            def gate(name, inp):
                w = pop[f'{name}.weight'].view(pop_size, -1)
                b = pop[f'{name}.bias'].view(pop_size)
                return heaviside(inp @ w.T + b)  # [num, pop]

            def and_gate(name, x, y):
                w = pop[f'{name}.weight'].view(pop_size, 2)
                b = pop[f'{name}.bias'].view(pop_size)
                inp = torch.stack([x, y], dim=-1)
                return heaviside((inp * w).sum(-1) + b)

            exp_zero = gate(f'{family}.classify.exp_zero', exp_in)
            exp_maxg = gate(f'{family}.classify.exp_max', exp_in)
            frac_zero = gate(f'{family}.classify.frac_zero', frac_in)
            frac_nz = gate(f'{family}.classify.frac_nonzero', frac_in)

            checks = [
                (f'{family}.classify.exp_zero', exp_zero,
                 [e == 0 for e, _ in cases]),
                (f'{family}.classify.exp_max', exp_maxg,
                 [e == exp_max_val for e, _ in cases]),
                (f'{family}.classify.frac_zero', frac_zero,
                 [f == 0 for _, f in cases]),
                (f'{family}.classify.frac_nonzero', frac_nz,
                 [f != 0 for _, f in cases]),
                (f'{family}.classify.is_zero.and',
                 and_gate(f'{family}.classify.is_zero.and', exp_zero, frac_zero),
                 [e == 0 and f == 0 for e, f in cases]),
                (f'{family}.classify.is_subnormal.and',
                 and_gate(f'{family}.classify.is_subnormal.and', exp_zero, frac_nz),
                 [e == 0 and f != 0 for e, f in cases]),
                (f'{family}.classify.is_inf.and',
                 and_gate(f'{family}.classify.is_inf.and', exp_maxg, frac_zero),
                 [e == exp_max_val and f == 0 for e, f in cases]),
                (f'{family}.classify.is_nan.and',
                 and_gate(f'{family}.classify.is_nan.and', exp_maxg, frac_nz),
                 [e == exp_max_val and f != 0 for e, f in cases]),
            ]
            for name, out, exp_list in checks:
                expected = torch.tensor([1.0 if x else 0.0 for x in exp_list],
                                        device=self.device)
                correct = (out == expected.unsqueeze(1)).float().sum(0)
                scores += correct
                total += num
                self._record(name, int(correct[0].item()), num, [])
                if debug:
                    r = self.results[-1]
                    print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except KeyError as e:
            if debug:
                print(f"  {family}.classify: SKIP (missing {e})")

        return scores, total

    def _test_float_cmp_composed(self, pop: Dict, family: str, exp_bits: int,
                                 frac_bits: int, debug: bool) -> Tuple[torch.Tensor, int]:
        """Composed IEEE comparison test driven entirely by the shipped
        wiring: the netlist is reconstructed from the .inputs metadata and
        evaluated end to end, then checked against exact IEEE semantics
        (NaN unordered, +0 == -0, subnormals ordered, mixed signs)."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print(f"\n=== {family.upper()} CMP (composed, from .inputs netlist) ===")

        prefix = f"{family}.cmp"
        try:
            ne = NetlistEvaluator(pop, self.signal_registry, prefix, pop_size=pop_size)
        except (KeyError, ValueError) as e:
            if debug:
                print(f"  {prefix} composed: SKIP ({e})")
            return scores, 0
        if f"{prefix}.same_sign" not in ne.gates:
            if debug:
                print(f"  {prefix} composed: SKIP (pre-composition wiring)")
            return scores, 0

        directed, randoms = float_test_words(exp_bits, frac_bits)
        pairs = [(x, y) for x in directed for y in directed]
        pairs += list(zip(randoms, randoms[1:]))
        pairs += [(r, r) for r in randoms[:8]]

        W = 1 + exp_bits + frac_bits
        a_words = torch.tensor([p[0] for p in pairs], dtype=torch.long)
        b_words = torch.tensor([p[1] for p in pairs], dtype=torch.long)
        ext = {}
        for i in range(W):
            ext[f"$a[{i}]"] = ((a_words >> (W - 1 - i)) & 1).float()
            ext[f"$b[{i}]"] = ((b_words >> (W - 1 - i)) & 1).float()
        try:
            out = ne.run(ext)
        except KeyError as e:
            if debug:
                print(f"  {prefix} composed: SKIP (unbound signal {e})")
            return scores, 0

        av = [float_bits_to_value(p[0], exp_bits, frac_bits) for p in pairs]
        bv = [float_bits_to_value(p[1], exp_bits, frac_bits) for p in pairs]
        ops = [
            ("eq", lambda x, y: x == y),
            ("lt", lambda x, y: x < y),
            ("gt", lambda x, y: x > y),
            ("le", lambda x, y: x <= y),
            ("ge", lambda x, y: x >= y),
        ]
        for op, fn in ops:
            expected = torch.tensor([1.0 if fn(x, y) else 0.0 for x, y in zip(av, bv)],
                                    device=self.device)
            got = out[f"{prefix}.{op}.result"]
            correct = (got == expected.unsqueeze(1)).float().sum(0)
            scores += correct
            total += len(pairs)
            failures = []
            if pop_size == 1:
                for i in range(len(pairs)):
                    if got[i, 0].item() != expected[i].item():
                        failures.append((list(pairs[i]), expected[i].item(), got[i, 0].item()))
            self._record(f"{prefix}.{op}.composed", int(correct[0].item()), len(pairs), failures[:10])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
                for inp, exp_v, got_v in (failures or [])[:4]:
                    print(f"    FAIL: a={inp[0]:#x} b={inp[1]:#x} expected {exp_v}, got {got_v}")

        return scores, total

    def _test_float_arith_composed(self, pop: Dict, family: str, op: str,
                                   oracle, exp_bits: int, frac_bits: int,
                                   debug: bool) -> Tuple[torch.Tensor, int]:
        """Composed arithmetic test for a self-contained float pipeline: the
        netlist is rebuilt from the .inputs metadata, evaluated end to end
        over IEEE edge-case and random operand pairs, and the assembled
        output word is compared to the exact integer oracle."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print(f"\n=== {family.upper()} {op.upper()} (composed, from .inputs netlist) ===")

        prefix = f"{family}.{op}"
        try:
            ne = NetlistEvaluator(pop, self.signal_registry, prefix, pop_size=pop_size)
        except (KeyError, ValueError) as e:
            if debug:
                print(f"  {prefix} composed: SKIP ({e})")
            return scores, 0
        if (f"{prefix}.sel.norm" not in ne.gates
                and f"{prefix}.sel.dp_norm" not in ne.gates):
            if debug:
                print(f"  {prefix} composed: SKIP (pre-composition wiring)")
            return scores, 0

        E, F = exp_bits, frac_bits
        W = 1 + E + F
        directed, randoms = float_test_words(E, F)
        pairs = [(x, y) for x in directed for y in directed]
        pairs += list(zip(randoms, randoms[1:]))
        pairs += [(r, r) for r in randoms[:8]]

        a_words = torch.tensor([p[0] for p in pairs], dtype=torch.long)
        b_words = torch.tensor([p[1] for p in pairs], dtype=torch.long)
        ext = {}
        for i in range(W):
            ext[f"$a[{i}]"] = ((a_words >> (W - 1 - i)) & 1).float()
            ext[f"$b[{i}]"] = ((b_words >> (W - 1 - i)) & 1).float()
        try:
            out = ne.run(ext)
        except KeyError as e:
            if debug:
                print(f"  {prefix} composed: SKIP (unbound signal {e})")
            return scores, 0

        # Assemble the output word: sign, exponent (LSB-first gate index),
        # fraction (LSB-first gate index).
        got = out[f"{prefix}.sign_out"].double() * float(1 << (E + F))
        for k in range(E):
            got = got + out[f"{prefix}.exp_out.bit{k}"].double() * float(1 << (F + k))
        for k in range(F):
            got = got + out[f"{prefix}.frac_out.bit{k}"].double() * float(1 << k)

        expected = torch.tensor(
            [float(oracle(p[0], p[1], E, F)) for p in pairs], dtype=torch.float64
        ).unsqueeze(1)
        correct = (got == expected).float().sum(0)
        scores += correct
        total += len(pairs)

        failures = []
        if pop_size == 1:
            for i in range(len(pairs)):
                if got[i, 0].item() != expected[i, 0].item():
                    failures.append((list(pairs[i]), int(expected[i, 0].item()),
                                     int(got[i, 0].item())))
        self._record(f"{prefix}.composed", int(correct[0].item()), len(pairs), failures[:10])
        if debug:
            r = self.results[-1]
            print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
            for inp, exp_v, got_v in (failures or [])[:4]:
                print(f"    FAIL: a={inp[0]:#x} b={inp[1]:#x} expected {exp_v:#x}, got {got_v:#x}")

        return scores, total

    def _test_float_fma_composed(self, pop: Dict, family: str, exp_bits: int,
                                 frac_bits: int, debug: bool) -> Tuple[torch.Tensor, int]:
        """Composed fused-multiply-add test: rebuild the netlist from .inputs,
        evaluate round(a*b+c) end to end over edge and random triples, compare
        to the single-rounding oracle."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        prefix = f"{family}.fma"
        try:
            ne = NetlistEvaluator(pop, self.signal_registry, prefix, pop_size=pop_size)
        except (KeyError, ValueError) as e:
            if debug:
                print(f"  {prefix} composed: SKIP ({e})")
            return scores, 0
        if f"{prefix}.norm" not in ne.gates:
            if debug:
                print(f"  {prefix} composed: SKIP (pre-composition wiring)")
            return scores, 0

        E, F = exp_bits, frac_bits
        W = 1 + E + F
        directed, randoms = float_test_words(E, F)
        base = directed + randoms[:24]
        triples = [(base[i], base[(i * 7 + 3) % len(base)], base[(i * 13 + 5) % len(base)])
                   for i in range(len(base))]
        triples += list(zip(randoms, randoms[1:], randoms[2:]))
        a_words = torch.tensor([t[0] for t in triples], dtype=torch.long)
        b_words = torch.tensor([t[1] for t in triples], dtype=torch.long)
        c_words = torch.tensor([t[2] for t in triples], dtype=torch.long)
        ext = {}
        for i in range(W):
            ext[f"$a[{i}]"] = ((a_words >> (W - 1 - i)) & 1).float()
            ext[f"$b[{i}]"] = ((b_words >> (W - 1 - i)) & 1).float()
            ext[f"$c[{i}]"] = ((c_words >> (W - 1 - i)) & 1).float()
        try:
            out = ne.run(ext)
        except KeyError as e:
            if debug:
                print(f"  {prefix} composed: SKIP (unbound signal {e})")
            return scores, 0

        got = out[f"{prefix}.sign_out"].double() * float(1 << (E + F))
        for k in range(E):
            got = got + out[f"{prefix}.exp_out.bit{k}"].double() * float(1 << (F + k))
        for k in range(F):
            got = got + out[f"{prefix}.frac_out.bit{k}"].double() * float(1 << k)
        expected = torch.tensor(
            [float(float_fma_oracle(t[0], t[1], t[2], E, F)) for t in triples],
            dtype=torch.float64).unsqueeze(1)
        correct = (got == expected).float().sum(0)
        scores += correct
        failures = []
        if pop_size == 1:
            for i in range(len(triples)):
                if got[i, 0].item() != expected[i, 0].item():
                    failures.append((list(triples[i]), int(expected[i, 0].item()),
                                     int(got[i, 0].item())))
        self._record(f"{prefix}.composed", int(correct[0].item()), len(triples), failures[:10])
        if debug:
            r = self.results[-1]
            print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
            for inp, exp_v, got_v in (failures or [])[:4]:
                print(f"    FAIL: a={inp[0]:#x} b={inp[1]:#x} c={inp[2]:#x} expected {exp_v:#x}, got {got_v:#x}")
        return scores, len(triples)

    # =========================================================================
    # FLOAT16 STRUCTURE CHECKS
    # =========================================================================

    def _test_float16_core(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Structure (shape) checks for the float16 core gate inventory."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== FLOAT16 CORE (structure) ===")

        expected_gates = [
            ('float16.unpack.bit0.weight', (1,)),
            ('float16.classify.exp_zero.weight', (5,)),
            ('float16.classify.exp_max.weight', (5,)),
            ('float16.classify.frac_zero.weight', (10,)),
            ('float16.classify.is_zero.and.weight', (2,)),
            ('float16.classify.is_nan.and.weight', (2,)),
            ('float16.pack.bit0.weight', (1,)),
        ]

        for name, expected_shape in expected_gates:
            try:
                tensor = pop[name]
                actual_shape = tuple(tensor.shape[1:])
                if actual_shape == expected_shape:
                    scores += 1
                    self._record(name, 1, 1, [])
                else:
                    self._record(name, 0, 1, [(expected_shape, actual_shape)])
                total += 1
                if debug:
                    r = self.results[-1]
                    print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
            except KeyError:
                if debug:
                    print(f"  {name}: SKIP (not found)")

        return scores, total

    def _test_float16_add(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Structure (shape) checks for the float16 addition stage gates."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== FLOAT16 ADD (structure) ===")

        expected_gates = [
            ('float16.add.pl.gt.weight', (15,)),  # payload-cascade final OR
            ('float16.add.exp_diff.fa0.ha1.sum.layer1.or.weight', (2,)),
            ('float16.add.align.s0.bit0.not_sel.weight', (1,)),
            ('float16.add.sign_xor.layer1.or.weight', (2,)),
            ('float16.add.addp.fa0.ha1.sum.layer1.or.weight', (2,)),
            ('float16.add.subp.not_s.bit0.weight', (1,)),
            ('float16.add.res.bit0.or.weight', (2,)),
            ('float16.add.lzc.nz.weight', (14,)),
            ('float16.add.sel.dp_norm.weight', (3,)),
        ]

        for name, expected_shape in expected_gates:
            try:
                tensor = pop[name]
                actual_shape = tuple(tensor.shape[1:])
                if actual_shape == expected_shape:
                    scores += 1
                    self._record(name, 1, 1, [])
                else:
                    self._record(name, 0, 1, [(expected_shape, actual_shape)])
                total += 1
                if debug:
                    r = self.results[-1]
                    print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
            except KeyError:
                if debug:
                    print(f"  {name}: SKIP (not found)")

        return scores, total

    def _test_float16_mul(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Structure (shape) checks for the float16 multiplication stage gates."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== FLOAT16 MUL (structure) ===")

        expected_gates = [
            ('float16.mul.sign_xor.layer1.or.weight', (2,)),
            ('float16.mul.exp_add.fa0.ha1.sum.layer1.or.weight', (2,)),
            ('float16.mul.exp_r.fa0.ha1.sum.layer1.or.weight', (2,)),
            ('float16.mul.mant_mul.pp.a0b0.weight', (2,)),
            ('float16.mul.mant_mul.acc.s0.fa0.ha1.sum.layer1.or.weight', (2,)),
            ('float16.mul.norm.bit0.or.weight', (2,)),
            ('float16.mul.sel.norm.weight', (3,)),
        ]

        for name, expected_shape in expected_gates:
            try:
                tensor = pop[name]
                actual_shape = tuple(tensor.shape[1:])
                if actual_shape == expected_shape:
                    scores += 1
                    self._record(name, 1, 1, [])
                else:
                    self._record(name, 0, 1, [(expected_shape, actual_shape)])
                total += 1
                if debug:
                    r = self.results[-1]
                    print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
            except KeyError:
                if debug:
                    print(f"  {name}: SKIP (not found)")

        return scores, total

    def _test_float16_div(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Structure (shape) checks for the float16 division stage gates."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== FLOAT16 DIV (structure) ===")

        expected_gates = [
            ('float16.div.sign_xor.layer1.or.weight', (2,)),
            ('float16.div.exp_nb.bit0.weight', (1,)),
            ('float16.div.exp_r.fa0.ha1.sum.layer1.or.weight', (2,)),
            ('float16.div.mant_div.stage0.cmp.weight', (1,)),  # bit-cascaded GE = NOT(LT) buffer
            ('float16.div.mant_div.stage0.q.weight', (2,)),
            ('float16.div.mant_div.stage0.sub.not_d.bit0.weight', (1,)),
            ('float16.div.mant_div.stage0.mux.bit0.not_sel.weight', (1,)),
            ('float16.div.sel.norm.weight', (3,)),
        ]

        for name, expected_shape in expected_gates:
            try:
                tensor = pop[name]
                actual_shape = tuple(tensor.shape[1:])
                if actual_shape == expected_shape:
                    scores += 1
                    self._record(name, 1, 1, [])
                else:
                    self._record(name, 0, 1, [(expected_shape, actual_shape)])
                total += 1
                if debug:
                    r = self.results[-1]
                    print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
            except KeyError:
                if debug:
                    print(f"  {name}: SKIP (not found)")

        return scores, total

    def _test_float16_cmp(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Structure (shape) checks for the float16 comparison gate inventory."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== FLOAT16 CMP (structure) ===")

        expected_gates = [
            ('float16.cmp.a.exp_max.weight', (5,)),
            ('float16.cmp.a.frac_nz.weight', (10,)),
            ('float16.cmp.a.is_nan.weight', (2,)),
            ('float16.cmp.either_nan.weight', (2,)),
            ('float16.cmp.sign_xor.layer1.or.weight', (2,)),
            ('float16.cmp.both_zero.weight', (2,)),
            ('float16.cmp.mag_a_gt_b.weight', (15,)),  # bit-cascaded final OR over 15 bits
            ('float16.cmp.eq.result.weight', (2,)),
            ('float16.cmp.lt.result.weight', (3,)),
            ('float16.cmp.gt.result.weight', (3,)),
        ]

        for name, expected_shape in expected_gates:
            try:
                tensor = pop[name]
                actual_shape = tuple(tensor.shape[1:])
                if actual_shape == expected_shape:
                    scores += 1
                    self._record(name, 1, 1, [])
                else:
                    self._record(name, 0, 1, [(expected_shape, actual_shape)])
                total += 1
                if debug:
                    r = self.results[-1]
                    print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
            except KeyError:
                if debug:
                    print(f"  {name}: SKIP (not found)")

        return scores, total

    # =========================================================================
    # FLOAT32 TESTS
    # =========================================================================

    def _test_float32_core(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Structure (shape) checks for the float32 core gate inventory."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== FLOAT32 CORE (structure) ===")

        expected_gates = [
            ('float32.unpack.bit0.weight', (1,)),
            ('float32.classify.exp_zero.weight', (8,)),
            ('float32.classify.exp_max.weight', (8,)),
            ('float32.classify.frac_zero.weight', (23,)),
            ('float32.classify.is_zero.and.weight', (2,)),
            ('float32.classify.is_nan.and.weight', (2,)),
            ('float32.pack.bit0.weight', (1,)),
        ]

        for name, expected_shape in expected_gates:
            try:
                tensor = pop[name]
                actual_shape = tuple(tensor.shape[1:])
                if actual_shape == expected_shape:
                    scores += 1
                    self._record(name, 1, 1, [])
                else:
                    self._record(name, 0, 1, [(expected_shape, actual_shape)])
                total += 1
                if debug:
                    r = self.results[-1]
                    print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
            except KeyError:
                if debug:
                    print(f"  {name}: SKIP (not found)")

        return scores, total

    def _test_float32_add(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Structure (shape) checks for the float32 addition stage gates."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== FLOAT32 ADD (structure) ===")

        expected_gates = [
            ('float32.add.pl.gt.weight', (31,)),  # payload-cascade final OR
            ('float32.add.exp_diff.fa0.ha1.sum.layer1.or.weight', (2,)),
            ('float32.add.align.s0.bit0.not_sel.weight', (1,)),
            ('float32.add.sign_xor.layer1.or.weight', (2,)),
            ('float32.add.addp.fa0.ha1.sum.layer1.or.weight', (2,)),
            ('float32.add.subp.not_s.bit0.weight', (1,)),
            ('float32.add.res.bit0.or.weight', (2,)),
            ('float32.add.lzc.nz.weight', (27,)),
            ('float32.add.sel.dp_norm.weight', (3,)),
        ]

        for name, expected_shape in expected_gates:
            try:
                tensor = pop[name]
                actual_shape = tuple(tensor.shape[1:])
                if actual_shape == expected_shape:
                    scores += 1
                    self._record(name, 1, 1, [])
                else:
                    self._record(name, 0, 1, [(expected_shape, actual_shape)])
                total += 1
                if debug:
                    r = self.results[-1]
                    print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
            except KeyError:
                if debug:
                    print(f"  {name}: SKIP (not found)")

        return scores, total

    def _test_float32_mul(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Structure (shape) checks for the float32 multiplication stage gates."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== FLOAT32 MUL (structure) ===")

        expected_gates = [
            ('float32.mul.sign_xor.layer1.or.weight', (2,)),
            ('float32.mul.exp_add.fa0.ha1.sum.layer1.or.weight', (2,)),
            ('float32.mul.exp_r.fa0.ha1.sum.layer1.or.weight', (2,)),
            ('float32.mul.mant_mul.pp.a0b0.weight', (2,)),
            ('float32.mul.mant_mul.acc.s0.fa0.ha1.sum.layer1.or.weight', (2,)),
            ('float32.mul.norm.bit0.or.weight', (2,)),
            ('float32.mul.sel.norm.weight', (3,)),
        ]

        for name, expected_shape in expected_gates:
            try:
                tensor = pop[name]
                actual_shape = tuple(tensor.shape[1:])
                if actual_shape == expected_shape:
                    scores += 1
                    self._record(name, 1, 1, [])
                else:
                    self._record(name, 0, 1, [(expected_shape, actual_shape)])
                total += 1
                if debug:
                    r = self.results[-1]
                    print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
            except KeyError:
                if debug:
                    print(f"  {name}: SKIP (not found)")

        return scores, total

    def _test_float32_div(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Structure (shape) checks for the float32 division stage gates."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== FLOAT32 DIV (structure) ===")

        expected_gates = [
            ('float32.div.sign_xor.layer1.or.weight', (2,)),
            ('float32.div.exp_nb.bit0.weight', (1,)),
            ('float32.div.exp_r.fa0.ha1.sum.layer1.or.weight', (2,)),
            ('float32.div.mant_div.stage0.cmp.weight', (1,)),  # bit-cascaded GE = NOT(LT) buffer
            ('float32.div.mant_div.stage0.q.weight', (2,)),
            ('float32.div.mant_div.stage0.sub.not_d.bit0.weight', (1,)),
            ('float32.div.mant_div.stage0.mux.bit0.not_sel.weight', (1,)),
            ('float32.div.sel.norm.weight', (3,)),
        ]

        for name, expected_shape in expected_gates:
            try:
                tensor = pop[name]
                actual_shape = tuple(tensor.shape[1:])
                if actual_shape == expected_shape:
                    scores += 1
                    self._record(name, 1, 1, [])
                else:
                    self._record(name, 0, 1, [(expected_shape, actual_shape)])
                total += 1
                if debug:
                    r = self.results[-1]
                    print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
            except KeyError:
                if debug:
                    print(f"  {name}: SKIP (not found)")

        return scores, total

    def _test_float32_cmp(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Structure (shape) checks for the float32 comparison gate inventory."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== FLOAT32 CMP (structure) ===")

        expected_gates = [
            ('float32.cmp.a.exp_max.weight', (8,)),
            ('float32.cmp.a.frac_nz.weight', (23,)),
            ('float32.cmp.a.is_nan.weight', (2,)),
            ('float32.cmp.either_nan.weight', (2,)),
            ('float32.cmp.sign_xor.layer1.or.weight', (2,)),
            ('float32.cmp.both_zero.weight', (2,)),
            ('float32.cmp.mag_a_gt_b.weight', (31,)),  # bit-cascaded final OR over 31 bits
            ('float32.cmp.eq.result.weight', (2,)),
            ('float32.cmp.lt.result.weight', (3,)),
            ('float32.cmp.gt.result.weight', (3,)),
        ]

        for name, expected_shape in expected_gates:
            try:
                tensor = pop[name]
                actual_shape = tuple(tensor.shape[1:])
                if actual_shape == expected_shape:
                    scores += 1
                    self._record(name, 1, 1, [])
                else:
                    self._record(name, 0, 1, [(expected_shape, actual_shape)])
                total += 1
                if debug:
                    r = self.results[-1]
                    print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
            except KeyError:
                if debug:
                    print(f"  {name}: SKIP (not found)")

        return scores, total

    # =========================================================================
    # INTEGRATION TESTS (Multi-circuit chains)
    # =========================================================================

    def _pop_modN(self, pop: Dict, pop_size: int, val_bits: torch.Tensor,
                  modulus: int) -> torch.Tensor:
        """Drive the bit-cascade modular.mod{N} divisibility detector.

        Returns a (pop_size,) tensor: 1 iff the 8-bit value (MSB-first bits in
        val_bits) is divisible by ``modulus``. Walks the per-multiple match
        gates (modular.modN.eq.k{val}.bit{i}.match -> .all -> top-level OR).
        """
        ks = [k for k in range(256) if k % modulus == 0]
        alls = []
        for k in ks:
            matches = []
            for i in range(8):
                w = pop[f'modular.mod{modulus}.eq.k{k}.bit{i}.match.weight'].view(pop_size, 1)
                b = pop[f'modular.mod{modulus}.eq.k{k}.bit{i}.match.bias'].view(pop_size)
                matches.append(heaviside(val_bits[i] * w[:, 0] + b))
            all_inp = torch.stack(matches, dim=-1)
            w_all = pop[f'modular.mod{modulus}.eq.k{k}.all.weight'].view(pop_size, 8)
            b_all = pop[f'modular.mod{modulus}.eq.k{k}.all.bias'].view(pop_size)
            alls.append(heaviside((all_inp * w_all).sum(-1) + b_all))
        top_inp = torch.stack(alls, dim=-1)
        w_top = pop[f'modular.mod{modulus}.weight'].view(pop_size, len(ks))
        b_top = pop[f'modular.mod{modulus}.bias'].view(pop_size)
        return heaviside((top_inp * w_top).sum(-1) + b_top)

    def _pop_cmp8bit(self, pop: Dict, pop_size: int,
                     a_bits: torch.Tensor, b_bits: torch.Tensor,
                     kind: str) -> torch.Tensor:
        """Drive the bit-cascade comparator (cmp8bit) over a population.

        Returns a (pop_size,) tensor of heaviside outputs for the requested
        comparison kind ('gt' | 'lt' | 'eq'). Bit 0 is MSB.
        """
        def apply(name: str, inp: torch.Tensor, fan_in: int) -> torch.Tensor:
            w = pop[f'{name}.weight'].view(pop_size, fan_in)
            b = pop[f'{name}.bias'].view(pop_size)
            return heaviside((inp * w).sum(-1) + b)

        # Per-bit primitives.
        bit_gt, bit_lt, bit_eq = [], [], []
        for i in range(8):
            ab = torch.stack([a_bits[i], b_bits[i]])
            bit_gt.append(apply(f'arithmetic.cmp8bit.bit{i}.gt', ab, 2))
            bit_lt.append(apply(f'arithmetic.cmp8bit.bit{i}.lt', ab, 2))
            eq_and = apply(f'arithmetic.cmp8bit.bit{i}.eq.layer1.and', ab, 2)
            eq_nor = apply(f'arithmetic.cmp8bit.bit{i}.eq.layer1.nor', ab, 2)
            eq_in = torch.stack([eq_and, eq_nor], dim=-1)
            w = pop[f'arithmetic.cmp8bit.bit{i}.eq.weight'].view(pop_size, 2)
            b = pop[f'arithmetic.cmp8bit.bit{i}.eq.bias'].view(pop_size)
            bit_eq.append(heaviside((eq_in * w).sum(-1) + b))

        # Cascade.
        cas_gt = [bit_gt[0]]
        cas_lt = [bit_lt[0]]
        for i in range(1, 8):
            eq_pref_in = torch.stack(bit_eq[:i], dim=-1)
            w_pref = pop[f'arithmetic.cmp8bit.cascade.eq_prefix.bit{i}.weight'].view(pop_size, i)
            b_pref = pop[f'arithmetic.cmp8bit.cascade.eq_prefix.bit{i}.bias'].view(pop_size)
            eq_pref = heaviside((eq_pref_in * w_pref).sum(-1) + b_pref)
            cas_in = torch.stack([eq_pref, bit_gt[i]], dim=-1)
            w_g = pop[f'arithmetic.cmp8bit.cascade.gt.bit{i}.weight'].view(pop_size, 2)
            b_g = pop[f'arithmetic.cmp8bit.cascade.gt.bit{i}.bias'].view(pop_size)
            cas_gt.append(heaviside((cas_in * w_g).sum(-1) + b_g))
            cas_in_lt = torch.stack([eq_pref, bit_lt[i]], dim=-1)
            w_l = pop[f'arithmetic.cmp8bit.cascade.lt.bit{i}.weight'].view(pop_size, 2)
            b_l = pop[f'arithmetic.cmp8bit.cascade.lt.bit{i}.bias'].view(pop_size)
            cas_lt.append(heaviside((cas_in_lt * w_l).sum(-1) + b_l))

        if kind == 'gt':
            inp = torch.stack(cas_gt, dim=-1)
            return apply('arithmetic.greaterthan8bit', inp, 8)
        if kind == 'lt':
            inp = torch.stack(cas_lt, dim=-1)
            return apply('arithmetic.lessthan8bit', inp, 8)
        if kind == 'eq':
            inp = torch.stack(bit_eq, dim=-1)
            return apply('arithmetic.equality8bit', inp, 8)
        raise ValueError(kind)

    def _test_integration(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
        """Test complex operations that chain multiple circuit families."""
        pop_size = next(iter(pop.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total = 0

        if debug:
            print("\n=== INTEGRATION TESTS ===")

        # Test 1: ADD then compare (A + B > C?)
        # Uses: ripple carry adder + comparator
        try:
            op_scores = torch.zeros(pop_size, device=self.device)
            op_total = 0

            tests = [(10, 20, 25), (100, 50, 200), (255, 1, 0), (0, 0, 1)]
            for a, b, c in tests:
                sum_val = (a + b) & 0xFF
                expected = float(sum_val > c)

                # Compute sum bits
                sum_bits = torch.tensor([((sum_val >> (7 - i)) & 1) for i in range(8)],
                                       device=self.device, dtype=torch.float32)
                c_bits = torch.tensor([((c >> (7 - i)) & 1) for i in range(8)],
                                     device=self.device, dtype=torch.float32)

                # Drive sum_bits vs c_bits through the bit-cascade comparator.
                out = self._pop_cmp8bit(pop, pop_size, sum_bits, c_bits, 'gt')
                correct = (out == expected).float()
                op_scores += correct
                op_total += 1

            scores += op_scores
            total += op_total
            self._record('integration.add_then_compare', int(op_scores[0].item()), op_total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  integration.add_then_compare: SKIP ({e})")

        # Test 2: MUL then MOD (A * B mod 3 == 0?)
        # Uses: partial products + modular arithmetic concept
        try:
            op_scores = torch.zeros(pop_size, device=self.device)
            op_total = 0

            tests = [(3, 5), (4, 6), (7, 11), (9, 9)]
            for a, b in tests:
                product = (a * b) & 0xFF
                expected = float(product % 3 == 0)

                # Drive product bits through the bit-cascade mod3 detector;
                # output is 1 iff product is divisible by 3.
                prod_bits = torch.tensor([((product >> (7 - i)) & 1) for i in range(8)],
                                         device=self.device, dtype=torch.float32)
                out = self._pop_modN(pop, pop_size, prod_bits, 3)
                op_scores += (out == expected).float()
                op_total += 1

            scores += op_scores
            total += op_total
            self._record('integration.mul_then_mod', int(op_scores[0].item()), op_total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  integration.mul_then_mod: SKIP ({e})")

        # Test 3: Shift then AND (SHL(A) & B)
        # Uses: shift + bitwise AND
        try:
            op_scores = torch.zeros(pop_size, device=self.device)
            op_total = 0

            tests = [(0b10101010, 0b11110000), (0b00001111, 0b01010101), (0xFF, 0x0F)]
            for a, b in tests:
                shifted_a = (a << 1) & 0xFF
                expected = shifted_a & b

                a_bits = torch.tensor([((a >> (7 - i)) & 1) for i in range(8)],
                                     device=self.device, dtype=torch.float32)
                b_bits = torch.tensor([((b >> (7 - i)) & 1) for i in range(8)],
                                     device=self.device, dtype=torch.float32)

                # Apply SHL
                shifted_bits = []
                for bit in range(8):
                    w = pop[f'alu.alu8bit.shl.bit{bit}.weight'].view(pop_size)
                    bias = pop[f'alu.alu8bit.shl.bit{bit}.bias'].view(pop_size)
                    if bit < 7:
                        inp = a_bits[bit + 1].expand(pop_size)
                    else:
                        inp = torch.zeros(pop_size, device=self.device)
                    shifted_bits.append(heaviside(inp * w + bias))

                # Apply AND
                and_bits = []
                w_and = pop['alu.alu8bit.and.weight'].view(pop_size, 8, 2)
                b_and = pop['alu.alu8bit.and.bias'].view(pop_size, 8)
                for bit in range(8):
                    inp = torch.stack([shifted_bits[bit], b_bits[bit].expand(pop_size)], dim=-1)
                    and_bits.append(heaviside((inp * w_and[:, bit]).sum(-1) + b_and[:, bit]))

                out_val = torch.zeros(pop_size, device=self.device)
                for i in range(8):
                    out_val += and_bits[i] * (1 << (7 - i))
                op_scores += (out_val == expected).float()
                op_total += 1

            scores += op_scores
            total += op_total
            self._record('integration.shift_then_and', int(op_scores[0].item()), op_total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  integration.shift_then_and: SKIP ({e})")

        # Test 4: SUB then conditional (A - B, if result < 0 then NEG)
        # Uses: subtractor + comparator + conditional logic
        try:
            op_scores = torch.zeros(pop_size, device=self.device)
            op_total = 0

            tests = [(50, 30), (30, 50), (100, 100), (0, 1)]
            for a, b in tests:
                diff = (a - b) & 0xFF
                is_negative = a < b
                expected = (-diff & 0xFF) if is_negative else diff

                # Just verify the subtraction works correctly
                # (Full conditional logic would require control flow)
                a_bits = torch.tensor([((a >> (7 - i)) & 1) for i in range(8)],
                                     device=self.device, dtype=torch.float32)
                b_bits = torch.tensor([((b >> (7 - i)) & 1) for i in range(8)],
                                     device=self.device, dtype=torch.float32)

                # Drive a_bits vs b_bits through the bit-cascade LT comparator.
                lt_out = self._pop_cmp8bit(pop, pop_size, a_bits, b_bits, 'lt')

                op_scores += (lt_out == float(is_negative)).float()
                op_total += 1

            scores += op_scores
            total += op_total
            self._record('integration.sub_then_conditional', int(op_scores[0].item()), op_total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  integration.sub_then_conditional: SKIP ({e})")

        # Test 5: Complex expression: ((A + B) * 2) & 0xF0
        # Uses: adder + SHL + AND
        try:
            op_scores = torch.zeros(pop_size, device=self.device)
            op_total = 0

            tests = [(10, 20), (50, 50), (127, 1), (0, 0)]
            for a, b in tests:
                sum_val = (a + b) & 0xFF
                doubled = (sum_val << 1) & 0xFF
                expected = doubled & 0xF0

                sum_bits = torch.tensor([((sum_val >> (7 - i)) & 1) for i in range(8)],
                                       device=self.device, dtype=torch.float32)
                mask_bits = torch.tensor([1, 1, 1, 1, 0, 0, 0, 0],
                                        device=self.device, dtype=torch.float32)

                # Apply SHL
                shifted_bits = []
                for bit in range(8):
                    w = pop[f'alu.alu8bit.shl.bit{bit}.weight'].view(pop_size)
                    bias = pop[f'alu.alu8bit.shl.bit{bit}.bias'].view(pop_size)
                    if bit < 7:
                        inp = sum_bits[bit + 1].expand(pop_size)
                    else:
                        inp = torch.zeros(pop_size, device=self.device)
                    shifted_bits.append(heaviside(inp * w + bias))

                # Apply AND with mask
                w_and = pop['alu.alu8bit.and.weight'].view(pop_size, 8, 2)
                b_and = pop['alu.alu8bit.and.bias'].view(pop_size, 8)
                result_bits = []
                for bit in range(8):
                    inp = torch.stack([shifted_bits[bit], mask_bits[bit].expand(pop_size)], dim=-1)
                    result_bits.append(heaviside((inp * w_and[:, bit]).sum(-1) + b_and[:, bit]))

                out_val = torch.zeros(pop_size, device=self.device)
                for i in range(8):
                    out_val += result_bits[i] * (1 << (7 - i))
                op_scores += (out_val == expected).float()
                op_total += 1

            scores += op_scores
            total += op_total
            self._record('integration.complex_expr', int(op_scores[0].item()), op_total, [])
            if debug:
                r = self.results[-1]
                print(f"  {r.name}: {r.passed}/{r.total} {'PASS' if r.success else 'FAIL'}")
        except (KeyError, RuntimeError) as e:
            if debug:
                print(f"  integration.complex_expr: SKIP ({e})")

        return scores, total

    # =========================================================================
    # MAIN EVALUATE
    # =========================================================================

    def evaluate(self, population: Dict[str, torch.Tensor], debug: bool = False) -> torch.Tensor:
        """
        Evaluate population fitness with per-circuit reporting.

        Args:
            population: Dict of tensors, each with shape [pop_size, ...]
            debug: If True, print per-circuit results

        Returns:
            Tensor of fitness scores [pop_size], normalized to [0, 1]
        """
        self.results = []
        self.category_scores = {}

        pop_size = next(iter(population.values())).shape[0]
        scores = torch.zeros(pop_size, device=self.device)
        total_tests = 0

        # Boolean gates
        s, t = self._test_boolean_gates(population, debug)
        scores += s
        total_tests += t
        self.category_scores['boolean'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

        # Half adder
        s, t = self._test_halfadder(population, debug)
        scores += s
        total_tests += t
        self.category_scores['halfadder'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

        # Full adder
        s, t = self._test_fulladder(population, debug)
        scores += s
        total_tests += t
        self.category_scores['fulladder'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

        # Ripple carry adders
        for bits in [2, 4, 8]:
            s, t = self._test_ripplecarry(population, bits, debug)
            scores += s
            total_tests += t
            self.category_scores[f'ripplecarry{bits}'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

        # 16/32-bit circuits (if present)
        for bits in [16, 32]:
            if f'arithmetic.ripplecarry{bits}bit.fa0.ha1.sum.layer1.or.weight' in population:
                if debug:
                    print(f"\n{'=' * 60}")
                    print(f" {bits}-BIT CIRCUITS")
                    print(f"{'=' * 60}")

                s, t = self._test_ripplecarry(population, bits, debug)
                scores += s
                total_tests += t
                self.category_scores[f'ripplecarry{bits}'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

                s, t = self._test_comparators_nbits(population, bits, debug)
                scores += s
                total_tests += t
                self.category_scores[f'comparators{bits}'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

                if f'arithmetic.sub{bits}bit.not_b.bit0.weight' in population:
                    s, t = self._test_subtractor_nbits(population, bits, debug)
                    scores += s
                    total_tests += t
                    self.category_scores[f'subtractor{bits}'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

                if f'alu.alu{bits}bit.and.bit0.weight' in population:
                    s, t = self._test_bitwise_nbits(population, bits, debug)
                    scores += s
                    total_tests += t
                    self.category_scores[f'bitwise{bits}'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

                if f'alu.alu{bits}bit.shl.bit0.weight' in population:
                    s, t = self._test_shifts_nbits(population, bits, debug)
                    scores += s
                    total_tests += t
                    self.category_scores[f'shifts{bits}'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

                if f'alu.alu{bits}bit.inc.bit0.xor.layer1.or.weight' in population:
                    s, t = self._test_inc_dec_nbits(population, bits, debug)
                    scores += s
                    total_tests += t
                    self.category_scores[f'incdec{bits}'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

                if f'alu.alu{bits}bit.neg.not.bit0.weight' in population:
                    s, t = self._test_neg_nbits(population, bits, debug)
                    scores += s
                    total_tests += t
                    self.category_scores[f'neg{bits}'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

                if f'combinational.barrelshifter{bits}.layer0.bit0.not_sel.weight' in population:
                    s, t = self._test_barrel_shifter_nbits(population, bits, debug)
                    scores += s
                    total_tests += t
                    self.category_scores[f'barrelshifter{bits}'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

                if f'combinational.priorityencoder{bits}.valid.weight' in population:
                    s, t = self._test_priority_encoder_nbits(population, bits, debug)
                    scores += s
                    total_tests += t
                    self.category_scores[f'priorityencoder{bits}'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

        # 3-operand adder
        s, t = self._test_add3(population, debug)
        scores += s
        total_tests += t
        self.category_scores['add3'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

        # Order of operations (A + B × C)
        s, t = self._test_expr_add_mul(population, debug)
        scores += s
        total_tests += t
        self.category_scores['expr_add_mul'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

        # Comparators
        s, t = self._test_comparators(population, debug)
        scores += s
        total_tests += t
        self.category_scores['comparators'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

        # Threshold gates
        s, t = self._test_threshold_gates(population, debug)
        scores += s
        total_tests += t
        self.category_scores['threshold'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

        # Modular arithmetic
        s, t = self._test_modular_all(population, debug)
        scores += s
        total_tests += t
        self.category_scores['modular'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

        # Pattern recognition
        s, t = self._test_patterns(population, debug)
        scores += s
        total_tests += t
        self.category_scores['patterns'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

        # Error detection
        s, t = self._test_error_detection(population, debug)
        scores += s
        total_tests += t
        self.category_scores['error_detection'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

        # Combinational
        s, t = self._test_combinational(population, debug)
        scores += s
        total_tests += t
        self.category_scores['combinational'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

        # Control flow
        s, t = self._test_control_flow(population, debug)
        scores += s
        total_tests += t
        self.category_scores['control'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

        # ALU
        s, t = self._test_alu_ops(population, debug)
        scores += s
        total_tests += t
        self.category_scores['alu'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

        # Manifest
        s, t = self._test_manifest(population, debug)
        scores += s
        total_tests += t
        self.category_scores['manifest'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

        # Memory
        s, t = self._test_memory(population, debug)
        scores += s
        total_tests += t
        self.category_scores['memory'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

        # Float16 circuits (if present)
        if 'float16.unpack.bit0.weight' in population:
            if debug:
                print(f"\n{'=' * 60}")
                print(f" FLOAT16 CIRCUITS")
                print(f"{'=' * 60}")

            s, t = self._test_float_unpack_pack(population, 'float16', 16, debug)
            scores += s
            total_tests += t
            self.category_scores['float16_unpack_pack'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            s, t = self._test_float_classify(population, 'float16', 5, 10, debug)
            scores += s
            total_tests += t
            self.category_scores['float16_classify'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            s, t = self._test_float_cmp_composed(population, 'float16', 5, 10, debug)
            scores += s
            total_tests += t
            self.category_scores['float16_cmp_composed'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            s, t = self._test_float_arith_composed(population, 'float16', 'mul',
                                                   float_mul_oracle, 5, 10, debug)
            scores += s
            total_tests += t
            self.category_scores['float16_mul_composed'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            s, t = self._test_float_arith_composed(population, 'float16', 'div',
                                                   float_div_oracle, 5, 10, debug)
            scores += s
            total_tests += t
            self.category_scores['float16_div_composed'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            s, t = self._test_float_arith_composed(population, 'float16', 'add',
                                                   float_add_oracle, 5, 10, debug)
            scores += s
            total_tests += t
            self.category_scores['float16_add_composed'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            if 'float16.fma.norm.weight' in population:
                s, t = self._test_float_fma_composed(population, 'float16', 5, 10, debug)
                scores += s
                total_tests += t
                self.category_scores['float16_fma_composed'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            s, t = self._test_float16_core(population, debug)
            scores += s
            total_tests += t
            self.category_scores['float16_core'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            if 'float16.add.pl.gt.weight' in population:
                s, t = self._test_float16_add(population, debug)
                scores += s
                total_tests += t
                self.category_scores['float16_add'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            if 'float16.mul.sign_xor.layer1.or.weight' in population:
                s, t = self._test_float16_mul(population, debug)
                scores += s
                total_tests += t
                self.category_scores['float16_mul'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            if 'float16.div.sign_xor.layer1.or.weight' in population:
                s, t = self._test_float16_div(population, debug)
                scores += s
                total_tests += t
                self.category_scores['float16_div'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            if 'float16.cmp.a.exp_max.weight' in population:
                s, t = self._test_float16_cmp(population, debug)
                scores += s
                total_tests += t
                self.category_scores['float16_cmp'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

        # Float32 circuits (if present)
        if 'float32.unpack.bit0.weight' in population:
            if debug:
                print(f"\n{'=' * 60}")
                print(f" FLOAT32 CIRCUITS")
                print(f"{'=' * 60}")

            s, t = self._test_float_unpack_pack(population, 'float32', 32, debug)
            scores += s
            total_tests += t
            self.category_scores['float32_unpack_pack'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            s, t = self._test_float_classify(population, 'float32', 8, 23, debug)
            scores += s
            total_tests += t
            self.category_scores['float32_classify'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            s, t = self._test_float_cmp_composed(population, 'float32', 8, 23, debug)
            scores += s
            total_tests += t
            self.category_scores['float32_cmp_composed'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            s, t = self._test_float_arith_composed(population, 'float32', 'mul',
                                                   float_mul_oracle, 8, 23, debug)
            scores += s
            total_tests += t
            self.category_scores['float32_mul_composed'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            s, t = self._test_float_arith_composed(population, 'float32', 'div',
                                                   float_div_oracle, 8, 23, debug)
            scores += s
            total_tests += t
            self.category_scores['float32_div_composed'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            s, t = self._test_float_arith_composed(population, 'float32', 'add',
                                                   float_add_oracle, 8, 23, debug)
            scores += s
            total_tests += t
            self.category_scores['float32_add_composed'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            if 'float32.fma.norm.weight' in population:
                s, t = self._test_float_fma_composed(population, 'float32', 8, 23, debug)
                scores += s
                total_tests += t
                self.category_scores['float32_fma_composed'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            s, t = self._test_float32_core(population, debug)
            scores += s
            total_tests += t
            self.category_scores['float32_core'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            if 'float32.add.pl.gt.weight' in population:
                s, t = self._test_float32_add(population, debug)
                scores += s
                total_tests += t
                self.category_scores['float32_add'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            if 'float32.mul.sign_xor.layer1.or.weight' in population:
                s, t = self._test_float32_mul(population, debug)
                scores += s
                total_tests += t
                self.category_scores['float32_mul'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            if 'float32.div.sign_xor.layer1.or.weight' in population:
                s, t = self._test_float32_div(population, debug)
                scores += s
                total_tests += t
                self.category_scores['float32_div'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

            if 'float32.cmp.a.exp_max.weight' in population:
                s, t = self._test_float32_cmp(population, debug)
                scores += s
                total_tests += t
                self.category_scores['float32_cmp'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

        # Cross-family integration tests (chain ripple-carry, comparator,
        # modular, shifts, subtractor). Each test is internally guarded with
        # try/except, so unsupported variants silently skip individual tests.
        if 'arithmetic.cmp8bit.bit0.gt.weight' in population:
            s, t = self._test_integration(population, debug)
            scores += s
            total_tests += t
            self.category_scores['integration'] = (s[0].item() if pop_size == 1 else s.mean().item(), t)

        self.total_tests = total_tests

        if debug:
            print("\n" + "=" * 60)
            print("CATEGORY SUMMARY")
            print("=" * 60)
            for cat, (got, expected) in sorted(self.category_scores.items()):
                pct = 100 * got / expected if expected > 0 else 0
                status = "PASS" if got == expected else "FAIL"
                print(f"  {cat:20} {int(got):6}/{expected:6} ({pct:6.2f}%) [{status}]")

            print("\n" + "=" * 60)
            print("CIRCUIT FAILURES")
            print("=" * 60)
            failed = [r for r in self.results if not r.success]
            if failed:
                for r in failed[:20]:
                    print(f"  {r.name}: {r.passed}/{r.total}")
                    if r.failures:
                        print(f"    First failure: {r.failures[0]}")
                if len(failed) > 20:
                    print(f"  ... and {len(failed) - 20} more")
            else:
                print("  None!")

        return scores / total_tests if total_tests > 0 else scores


def main():
    parser = argparse.ArgumentParser(description='Unified Evaluation Suite for 8-bit Threshold Computer')
    parser.add_argument('--model', type=str, default=MODEL_PATH, help='Path to safetensors model')
    parser.add_argument('--device', type=str, default='cuda', help='Device: cuda or cpu')
    parser.add_argument('--pop_size', type=int, default=1, help='Population size for batched evaluation')
    parser.add_argument('--quiet', action='store_true', help='Suppress detailed output')
    parser.add_argument('--cpu-test', action='store_true', help='Run CPU smoke test (LOAD, ADD, STORE, HALT)')
    args = parser.parse_args()

    if args.cpu_test:
        return run_smoke_test()

    print("=" * 70)
    print(" UNIFIED EVALUATION SUITE")
    print("=" * 70)

    print(f"\nLoading model from {args.model}...")
    model = load_model(args.model)
    print(f"  Loaded {len(model)} tensors, {sum(t.numel() for t in model.values()):,} params")

    print(f"\nInitializing evaluator on {args.device}...")
    evaluator = BatchedFitnessEvaluator(device=args.device, model_path=args.model)

    print(f"\nCreating population (size {args.pop_size})...")
    population = create_population(model, pop_size=args.pop_size, device=args.device)

    print("\nRunning evaluation...")
    if args.device == 'cuda':
        torch.cuda.synchronize()
    start = time.perf_counter()

    fitness = evaluator.evaluate(population, debug=not args.quiet)

    if args.device == 'cuda':
        torch.cuda.synchronize()
    elapsed = time.perf_counter() - start

    print("\n" + "=" * 70)
    print("RESULTS")
    print("=" * 70)

    if args.pop_size == 1:
        print(f"  Fitness: {fitness[0].item():.6f}")
    else:
        print(f"  Mean Fitness: {fitness.mean().item():.6f}")
        print(f"  Min Fitness:  {fitness.min().item():.6f}")
        print(f"  Max Fitness:  {fitness.max().item():.6f}")

    print(f"  Total tests: {evaluator.total_tests}")
    print(f"  Time: {elapsed * 1000:.2f} ms")

    if args.pop_size > 1:
        print(f"  Throughput: {args.pop_size / elapsed:.0f} evals/sec")
        perfect = (fitness >= 0.9999).sum().item()
        print(f"  Perfect (>=99.99%): {perfect}/{args.pop_size}")

    if fitness[0].item() >= 0.9999:
        print("\n  STATUS: PASS")
        return 0
    else:
        failed_count = int((1 - fitness[0].item()) * evaluator.total_tests)
        print(f"\n  STATUS: FAIL ({failed_count} tests failed)")
        return 1


if __name__ == '__main__':
    exit(main())