"""Reversible threshold computer.

A conventional processor's state transition is not injective: overwriting a
register or a carry destroys information, which by Landauer's principle sets a
floor of kT ln 2 of dissipated energy per erased bit. This machine is built so
that the whole state transition T is a bijection, hence realizable with no
logical erasure and (on an adiabatically driven substrate) no Landauer floor.

Everything is expressed in the repository's threshold substrate. The reversible
primitives are threshold circuits whose input->output map happens to be a
permutation:

    NOT   t              -> t' = 1 - t
    CNOT  c t            -> t' = t XOR c
    TOFF  a b t          -> t' = t XOR (a AND b)          (Toffoli, universal)
    FRED  c x y          -> (x,y)' = (c?y:x, c?x:y)       (Fredkin, controlled swap)

Each target update is XOR(target, product-of-controls); XOR and AND are the same
Heaviside threshold gates used everywhere else in the repo, so a reversible gate
is a small threshold network that is bijective on its wires, and a composition of
them is a bijective threshold network on the register.

This module builds up from those gates to a reversible in-place adder (Cuccaro),
and later files add the reversible ALU, ISA, and the bijective machine step.
Reversibility is not asserted, it is verified: every construction is checked to
be a permutation of its state space, exhaustively at small widths.
"""
from __future__ import annotations
from typing import List

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


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


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


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


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


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


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


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


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


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


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


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


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


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


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


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


def incr(reg, b_bits, one_bits, carry):
    """b <- b + 1 mod 2^n. `one_bits` is a register holding the constant 1
    (LSB set), restored on exit; `carry` is a clean ancilla, restored."""
    add_into(reg, one_bits, b_bits, carry)


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


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


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


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


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


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


def bennett(reg, a_b, b_b, c_b, scr_b, out_b, carry):
    """Bennett compute-copy-uncompute for the irreversible f(a,b,c)=(a+b) XOR c.
    Maps (a,b,c,0,0) -> (a,b,c,0,f) with inputs preserved and scratch cleaned, so
    a function that discards information as a standalone map is realized by a
    reversible circuit. The op list is returned so the inverse is the reverse."""
    ops = []
    # compute f into scratch (scratch starts 0): scratch += a; scratch += b; scratch ^= c
    ops += _adder_ops(a_b, scr_b, carry)
    ops += _adder_ops(b_b, scr_b, carry)
    ops += [(CNOT, c, s) for c, s in zip(c_b, scr_b)]
    # copy scratch -> out
    ops += [(CNOT, s, o) for s, o in zip(scr_b, out_b)]
    # uncompute scratch (reverse of the compute prefix)
    n_comp = len(_adder_ops(a_b, scr_b, carry)) * 2 + len(c_b)
    ops += [(g, *args) for g, *args in reversed(ops[:n_comp])]
    _apply(reg, ops)
    return ops


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


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

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

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

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


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