"""Structured reversible programs over the reversible ALU.

The flat instruction machine (reversible_cpu.py) proves the per-step transition
is a bijection. This file is the structured layer that makes real programs
convenient and their reversibility obvious: statements are reversible register
updates and a bounded loop whose count register is read-only, so inverting a
program is a purely syntactic transform (reverse the statement order and invert
each statement) and running the inverse recovers the input exactly.

    ADD d,s   : d += s          inverse SUB d,s
    SUB d,s   : d -= s          inverse ADD d,s
    ADDI d,k  : d += k          inverse ADDI d,-k
    XOR d,s   : d ^= s          self-inverse
    NEG d     : d = -d          self-inverse
    SWAP a,b  : a,b = b,a       self-inverse
    FOR c B   : run B, c times  inverse FOR c invert(B)   (B must not write c)
    IF p T E  : Janus conditional with an exit assertion (reversible)

Every update is the value-level image of the verified reversible threshold
circuits in reversible.py; a FOR whose count register is preserved is a
reversible loop because both directions read the same count.
"""
from __future__ import annotations
from typing import Dict, List, Tuple

WIDTH = 8
MASK = (1 << WIDTH) - 1


def _regs_written(stmts) -> set:
    w = set()
    for st in stmts:
        op = st[0]
        if op in ("ADD", "SUB", "ADDI", "XOR", "NEG"):
            w.add(st[1])
        elif op == "SWAP":
            w.add(st[1]); w.add(st[2])
        elif op == "FOR":
            w |= _regs_written(st[2])
        elif op == "IF":
            w |= _regs_written(st[2]) | _regs_written(st[3])
    return w


def invert(stmts: List[tuple]) -> List[tuple]:
    out = []
    for st in reversed(stmts):
        op = st[0]
        if op == "ADD":
            out.append(("SUB", st[1], st[2]))
        elif op == "SUB":
            out.append(("ADD", st[1], st[2]))
        elif op == "ADDI":
            out.append(("ADDI", st[1], -st[2]))
        elif op in ("XOR", "NEG", "SWAP"):
            out.append(st)
        elif op == "FOR":
            out.append(("FOR", st[1], invert(st[2])))
        elif op == "IF":
            # Janus: reverse swaps predicate and exit assertion, inverts branches
            out.append(("IF", st[4], invert(st[2]), invert(st[3]), st[1]))
    return out


def run(stmts: List[tuple], s: Dict[str, int]) -> Dict[str, int]:
    for st in stmts:
        op = st[0]
        if op == "ADD":
            s[st[1]] = (s[st[1]] + s[st[2]]) & MASK
        elif op == "SUB":
            s[st[1]] = (s[st[1]] - s[st[2]]) & MASK
        elif op == "ADDI":
            s[st[1]] = (s[st[1]] + st[2]) & MASK
        elif op == "XOR":
            s[st[1]] ^= s[st[2]]
        elif op == "NEG":
            s[st[1]] = (-s[st[1]]) & MASK
        elif op == "SWAP":
            s[st[1]], s[st[2]] = s[st[2]], s[st[1]]
        elif op == "FOR":
            cnt, body = st[1], st[2]
            if cnt in _regs_written(body):
                raise ValueError("FOR count register must be read-only (irreversible otherwise)")
            for _ in range(s[cnt]):
                run(body, s)
        elif op == "IF":
            pred, then, els, exit_assert = st[1], st[2], st[3], st[4]
            if pred(s):
                run(then, s)
                assert exit_assert(s), "exit assertion violated (not reversible)"
            else:
                run(els, s)
                assert not exit_assert(s), "exit assertion violated (not reversible)"
    return s


# --- demonstration programs ---
MULTIPLY = [("FOR", "b", [("ADD", "acc", "a")])]          # acc += a, b times; a,b preserved
FIB = [("FOR", "n", [("ADD", "a", "b"), ("SWAP", "a", "b")])]   # (a,b)->(b,a+b), n times


def _test():
    ok = True

    # reversible multiply: acc = a*b, inputs preserved; inverse clears acc
    bad = 0
    for a in range(16):
        for b in range(16):
            s = {"a": a, "b": b, "acc": 0}
            run(MULTIPLY, s)
            if s["acc"] != (a * b) & MASK or s["a"] != a or s["b"] != b:
                bad += 1
            run(invert(MULTIPLY), s)                       # run backward
            if s != {"a": a, "b": b, "acc": 0}:
                bad += 1
    print(f"  reversible multiply acc=a*b, inputs preserved, inverse clears acc: "
          f"{'OK' if bad == 0 else f'FAIL({bad})'}")
    ok &= bad == 0

    # reversible Fibonacci: n steps forward, inverse recovers the seed
    bad = 0
    for n in range(12):
        s = {"a": 0, "b": 1, "n": n}
        run(FIB, s)
        # forward value check against a plain reference
        ra, rb = 0, 1
        for _ in range(n):
            ra, rb = rb, (ra + rb) & MASK
        if (s["a"], s["b"]) != (ra, rb):
            bad += 1
        run(invert(FIB), s)
        if s != {"a": 0, "b": 1, "n": n}:
            bad += 1
    print(f"  reversible Fibonacci n steps, inverse recovers seed: "
          f"{'OK' if bad == 0 else f'FAIL({bad})'}")
    ok &= bad == 0

    # a reversible conditional (Janus IF): swap when the operands differ. The
    # exit assertion a != b is true exactly when the then-branch ran (swapping
    # distinct values keeps them distinct; equal values are skipped), so the
    # reverse picks the right branch.
    prog = [("IF", lambda s: s["a"] != s["b"], [("SWAP", "a", "b")], [],
             lambda s: s["a"] != s["b"])]
    bad = 0
    for a in range(16):
        for b in range(16):
            s = {"a": a, "b": b}
            run(prog, s)
            if sorted([s["a"], s["b"]]) != sorted([a, b]):
                bad += 1
            run(invert(prog), s)
            if s != {"a": a, "b": b}:
                bad += 1
    print(f"  reversible conditional (Janus IF with exit assertion): "
          f"{'OK' if bad == 0 else f'FAIL({bad})'}")
    ok &= bad == 0
    return ok


if __name__ == "__main__":
    print("Reversible structured programs")
    print("PASS" if _test() else "FAIL")