"""Reversible register machine with a bijective state transition.

State s = (R[0..K-1], PC, BR, MEM[0..M-1]); the program is a read-only array
fetched by PC. Every instruction is a reversible update, and control flow is
reversible through a branch register BR: the program counter advances by
`PC += dir*BR` each cycle (BR = 1 for sequential flow), and a branch toggles BR
by XOR-ing `offset ^ 1`, so a matched branch at the destination restores BR to 1.
Because BR carries the control state, the exact same machine run with dir = -1
retraces the computation and reconstructs the input, dissipating nothing.

Instruction inverses (used for dir = -1):
    ADD<->SUB, ADDI(k)<->ADDI(-k), XOR/XORI/NEG/TOFF/EXCH self-inverse,
    ROL(k)<->ROL(-k); BRA/BEZ branch toggles are self-inverse.

The word-level updates are the reversible threshold circuits verified in
reversible.py (Cuccaro adder, bitwise Toffoli, rotate); this file is the
value-level machine whose single-step transition those circuits implement.
"""
from __future__ import annotations
from typing import Dict, List, Tuple, Optional


class RCPU:
    def __init__(self, program: List[tuple], k_regs=4, width=8, mem_words=16):
        self.prog = program
        self.K = k_regs
        self.W = width
        self.M = mem_words
        self.mask = (1 << width) - 1
        self.L = len(program)

    def new_state(self, regs=None, mem=None) -> dict:
        return {"R": list(regs) + [0] * (self.K - len(regs)) if regs else [0] * self.K,
                "PC": 0, "BR": 1, "MEM": list(mem) + [0] * (self.M - len(mem)) if mem else [0] * self.M}

    # ---- reversible instruction effects (forward and inverse) ----
    def _data(self, s, I, inverse: bool):
        R, MEM, m = s["R"], s["MEM"], self.mask
        op = I[0]
        if op == "ADD":
            d, r = I[1], I[2]
            R[d] = (R[d] - R[r]) & m if inverse else (R[d] + R[r]) & m
        elif op == "SUB":
            d, r = I[1], I[2]
            R[d] = (R[d] + R[r]) & m if inverse else (R[d] - R[r]) & m
        elif op == "ADDI":
            d, k = I[1], I[2]
            R[d] = (R[d] - k) & m if inverse else (R[d] + k) & m
        elif op == "XOR":
            R[I[1]] ^= R[I[2]]
        elif op == "XORI":
            R[I[1]] ^= (I[2] & m)
        elif op == "NEG":
            R[I[1]] = (-R[I[1]]) & m
        elif op == "TOFF":
            R[I[1]] ^= (R[I[2]] & R[I[3]])
        elif op == "ROL":
            k = (-I[2] if inverse else I[2]) % self.W
            R[I[1]] = ((R[I[1]] << k) | (R[I[1]] >> (self.W - k))) & m if k else R[I[1]]
        elif op == "EXCH":
            d, r = I[1], I[2]
            a = R[r] % self.M
            R[d], MEM[a] = MEM[a], R[d]
        # BRA/BEZ/HALT have no data effect

    def _toggle(self, s, I):
        """Reversible control: toggle BR for a taken branch (self-inverse)."""
        op = I[0]
        if op == "BRA":
            s["BR"] ^= (I[1] ^ 1)
        elif op == "BEZ":
            if s["R"][I[1]] == 0:
                s["BR"] ^= (I[2] ^ 1)

    # ---- single-step transition and its inverse ----
    def step(self, s):
        I = self.prog[s["PC"]]
        self._data(s, I, inverse=False)
        self._toggle(s, I)
        s["PC"] = (s["PC"] + s["BR"]) % self.L

    def step_back(self, s):
        s["PC"] = (s["PC"] - s["BR"]) % self.L
        I = self.prog[s["PC"]]
        self._toggle(s, I)               # self-inverse: restores BR
        self._data(s, I, inverse=True)

    def run(self, s, steps):
        for _ in range(steps):
            self.step(s)
        return s

    def run_back(self, s, steps):
        for _ in range(steps):
            self.step_back(s)
        return s


def _clone(s):
    return {"R": list(s["R"]), "PC": s["PC"], "BR": s["BR"], "MEM": list(s["MEM"])}


def _eq(a, b):
    return a["R"] == b["R"] and a["PC"] == b["PC"] and a["BR"] == b["BR"] and a["MEM"] == b["MEM"]


def test_straight_line():
    prog = [("ADD", 1, 0), ("XOR", 1, 0), ("NEG", 1), ("ADDI", 1, 7),
            ("TOFF", 2, 0, 1), ("ROL", 0, 1)]
    m = RCPU(prog, width=8)
    ok = True
    for a in (5, 0, 255, 100):
        for b in (3, 1, 200):
            s0 = m.new_state([a, b, 0, 0])
            s = _clone(s0)
            m.run(s, len(prog))
            fwd = _clone(s)
            m.run_back(s, len(prog))
            ok &= _eq(s, s0)               # round-trip recovers the exact input
    print(f"  straight-line round-trip (dir -1 recovers input): {'OK' if ok else 'FAIL'}")
    return ok


def test_bijection():
    """The machine's single-step transition is a bijection: step_back inverts
    step (and vice versa) for every instruction, over a sweep of states
    including the branch register. This is reversibility, program-independent."""
    import itertools
    W = 4
    prog = [
        ("ADD", 1, 0), ("SUB", 1, 0), ("ADDI", 2, 3), ("XOR", 1, 2),
        ("XORI", 0, 5), ("NEG", 3), ("TOFF", 3, 0, 1), ("ROL", 0, 1),
        ("EXCH", 2, 3), ("BRA", 2), ("BEZ", 1, 3), ("BEZ", 0, -2),
    ]
    m = RCPU(prog, k_regs=4, width=W, mem_words=4)
    bad = 0
    checked = 0
    rng = __import__("random").Random(0)
    for pc in range(len(prog)):
        for br in (1, 2, 3, -2, -1):
            for _ in range(60):
                s0 = {"R": [rng.randint(0, (1 << W) - 1) for _ in range(4)],
                      "PC": pc, "BR": br,
                      "MEM": [rng.randint(0, (1 << W) - 1) for _ in range(4)]}
                s = _clone(s0)
                m.step(s)
                m.step_back(s)
                if not _eq(s, s0):
                    bad += 1
                # and the other composition order
                s = _clone(s0)
                m.step_back(s)
                m.step(s)
                if not _eq(s, s0):
                    bad += 1
                checked += 2
    print(f"  step_back o step = id over {checked} (state, instruction) cases: "
          f"{'OK' if bad == 0 else f'FAIL({bad})'}")
    return bad == 0


if __name__ == "__main__":
    print("Reversible CPU")
    a = test_straight_line()
    b = test_bijection()
    print("PASS" if (a and b) else "FAIL")