"""Ship the reversible machine's arithmetic core as a threshold-gate artifact,
variants/neural_reversible.safetensors. The circuit is an in-place 8-bit adder
(b <- a+b) expressed as a sequence of reversible gates (CNOT, Toffoli); each gate
is realized by the Heaviside AND/XOR weights stored alongside, so the file holds
both the reversible program (the gate list) and the threshold substrate it runs
on. Round-trips the file and confirms the loaded circuit is a bijection that
computes b <- a+b with the addend and carry restored."""
from __future__ import annotations
import os
import random
import sys

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

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(ROOT, "src"))
import reversible as rv

OUT = os.path.join(ROOT, "variants", "neural_reversible.safetensors")
WIDTH = 8

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

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


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


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

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

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

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

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


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