"""Compile a reversible circuit to a ternary matrix stack and show the composed
transition is a permutation, realized on a crossbar with a measured noise
margin. This substantiates neural_reversible's no-erasure claim concretely with
the same matrix/crossbar machinery neural_matrix8 uses: a reversible circuit is
one product of ternary matrices with a Heaviside step, and because the circuit
is a bijection the composed map is a permutation of the state space, so every
matrix in the stack is applied without erasing information."""
from __future__ import annotations
import os
import sys

import torch

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


def adder_net(width: int):
    """The in-place Cuccaro adder (b <- a+b) as a feedforward ternary netlist:
    each reversible gate writes a fresh wire, so the input->output map is the
    permutation the circuit computes."""
    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)
    net = Net()
    inputs = [f"in{i}" for i in range(n)]
    cur = list(inputs)
    for k, (gate, *args) in enumerate(ops):
        if gate is rv.CNOT:
            c, t = args
            cur[t] = net.XOR(f"c{k}", cur[t], cur[c])          # t ^= c
        else:                                                   # Toffoli: t ^= a&b
            a, b, t = args
            tmp = net.AND(f"t{k}a", [cur[a], cur[b]])
            cur[t] = net.XOR(f"t{k}x", cur[t], tmp)
    return net, inputs, list(cur), n, a_bits, b_bits, carry


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


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


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


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

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

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


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