"""Reversible Margolus (partitioned) cellular automaton.

One fixed rule is applied to every 2x2 block, alternating the block partition
between even and odd alignment each step (the Margolus neighborhood). The state
is the lattice; there is no program counter, register file, or control circuit.

Rule (cells ordered TL,TR,BL,BR): rotate the block 180 degrees, except a pair of
particles on a diagonal (1001, 0110) swaps to the other diagonal. Both cases are
involutions and neither maps a state across the diagonal-pair boundary, so the
rule is a self-inverse permutation of the sixteen block states; the lattice
update is therefore a bijection and replaying the partition sequence in reverse
inverts the evolution.

Dynamics are the billiard-ball model's: isolated particles move ballistically on
diagonals and collisions deflect them reversibly. and_gate() computes AND from a
collision; ballistic transport plus this collision are the primitives of the
Fredkin-Toffoli universality construction (Margolus 1984).
"""
from __future__ import annotations
from typing import List, Tuple

Block = Tuple[int, int, int, int]


def rule(b: Block) -> Block:
    tl, tr, bl, br = b
    if b == (1, 0, 0, 1):
        return (0, 1, 1, 0)          # diagonal pair -> other diagonal (deflect)
    if b == (0, 1, 1, 0):
        return (1, 0, 0, 1)
    return (br, bl, tr, tl)          # otherwise rotate 180 degrees


def is_bijection() -> bool:
    outs = {rule(tuple((s >> k) & 1 for k in (3, 2, 1, 0))) for s in range(16)}
    return len(outs) == 16


def self_inverse() -> bool:
    return all(rule(rule(tuple((s >> k) & 1 for k in (3, 2, 1, 0))))
               == tuple((s >> k) & 1 for k in (3, 2, 1, 0)) for s in range(16))


def step(grid: List[List[int]], phase: int) -> List[List[int]]:
    """One Margolus update. phase 0 aligns blocks at even coordinates; phase 1
    offsets the partition by (1,1). Toroidal, so H and W must be even."""
    H, W = len(grid), len(grid[0])
    out = [row[:] for row in grid]
    o = phase
    for r0 in range(o, o + H, 2):
        for c0 in range(o, o + W, 2):
            r, r1 = r0 % H, (r0 + 1) % H
            c, c1 = c0 % W, (c0 + 1) % W
            nb = rule((grid[r][c], grid[r][c1], grid[r1][c], grid[r1][c1]))
            out[r][c], out[r][c1], out[r1][c], out[r1][c1] = nb
    return out


def run(grid: List[List[int]], nsteps: int, start_phase: int = 0) -> List[List[int]]:
    g = grid
    for n in range(nsteps):
        g = step(g, (start_phase + n) & 1)
    return g


def run_back(grid: List[List[int]], nsteps: int, start_phase: int = 0) -> List[List[int]]:
    """Undo `run`: replay the phase sequence in reverse; the rule is self-inverse."""
    phases = [(start_phase + n) & 1 for n in range(nsteps)]
    g = grid
    for p in reversed(phases):
        g = step(g, p)
    return g


# --- the block rule as Heaviside threshold gates ---
# rule(s) = rotate180(s) XOR is_diag(s) on every cell: rotation fixes diagonal
# pairs, and flipping all four cells of a rotated diagonal pair sends it to the
# other diagonal. is_diag detects the two diagonal-pair states.
def _H(x):
    return 1 if x >= 0 else 0


def _and(*xs):
    return _H(sum(xs) - len(xs))


def _or(*xs):
    return _H(sum(xs) - 1)


def _xor(a, b):
    return _and(_or(a, b), _H(1 - a - b))       # OR AND NAND, the family's XOR


def gate_rule(b: Block) -> Block:
    tl, tr, bl, br = b
    d = _or(_and(tl, 1 - tr, 1 - bl, br), _and(1 - tl, tr, bl, 1 - br))   # is_diag
    return (_xor(br, d), _xor(bl, d), _xor(tr, d), _xor(tl, d))


def _test_gates():
    ok = all(gate_rule(tuple((s >> k) & 1 for k in (3, 2, 1, 0)))
             == rule(tuple((s >> k) & 1 for k in (3, 2, 1, 0))) for s in range(16))
    print(f"  block rule as Heaviside threshold gates matches over 16 states: "
          f"{'OK' if ok else 'FAIL'}")
    return ok


# --- tests ---
def _rand_grid(H, W, seed):
    import random
    rng = random.Random(seed)
    return [[rng.randint(0, 1) for _ in range(W)] for _ in range(H)]


def _ball_positions(g):
    return {(r, c) for r, row in enumerate(g) for c, v in enumerate(row) if v}


def _test_rule():
    print(f"  block rule is a bijection of 16 states: {'OK' if is_bijection() else 'FAIL'}")
    print(f"  block rule is self-inverse: {'OK' if self_inverse() else 'FAIL'}")
    return is_bijection() and self_inverse()


def _test_reversibility():
    bad = 0
    for seed in range(20):
        g = _rand_grid(8, 8, seed)
        fwd = run(g, 25, start_phase=0)
        back = run_back(fwd, 25, start_phase=0)
        if back != g:
            bad += 1
    # particle count is conserved (the rule permutes cells within each block)
    g = _rand_grid(8, 8, 99)
    conserved = sum(sum(r) for r in g) == sum(sum(r) for r in run(g, 40))
    print(f"  lattice reversible (run then reverse recovers grid, 20 grids): "
          f"{'OK' if bad == 0 else f'FAIL({bad})'}")
    print(f"  particle number conserved: {'OK' if conserved else 'FAIL'}")
    return bad == 0 and conserved


def _test_ballistic():
    # a single particle travels in a straight diagonal line
    H = W = 16
    g = [[0] * W for _ in range(H)]
    g[2][2] = 1
    positions = [next(iter(_ball_positions(g)))]
    gg = g
    for n in range(8):
        gg = step(gg, n & 1)
        p = _ball_positions(gg)
        positions.append(next(iter(p)) if len(p) == 1 else None)
    ok = all(p is not None for p in positions)
    steady = ok and all(positions[i + 1] == (positions[i][0] + 1, positions[i][1] + 1)
                        for i in range(len(positions) - 1))
    print(f"  single particle stays a single particle: {'OK' if ok else 'FAIL'}")
    print(f"  and moves ballistically on the diagonal, +(1,1) per step: "
          f"{'OK' if steady else 'FAIL'}  trace={positions[:5]}")
    return ok and steady


def interaction_gate(a: int, b: int) -> dict:
    """One billiard-ball collision as a reversible interaction gate. Input
    particle A enters at (2,2) moving SE and B at (7,7) moving NW; at step 4
    three output cells carry A&B (the deflected paths), A&~B and ~A&B (the
    straight-through paths). AND plus routing by mirrors is functionally
    complete for the billiard-ball construction (Margolus 1984)."""
    H = W = 12
    g = [[0] * W for _ in range(H)]
    if a:
        g[2][2] = 1
    if b:
        g[7][7] = 1
    g = run(g, 4)
    return {"A_and_B": g[3][6], "A_and_notB": g[6][6], "notA_and_B": g[3][3]}


def and_gate(a: int, b: int) -> int:
    return interaction_gate(a, b)["A_and_B"]


def _test_gate():
    ok = True
    for a in (0, 1):
        for b in (0, 1):
            o = interaction_gate(a, b)
            ok &= (o["A_and_B"] == (a & b) and o["A_and_notB"] == (a & (1 - b))
                   and o["notA_and_B"] == ((1 - a) & b))
    print(f"  billiard-ball interaction gate (A&B, A&~B, ~A&B) over all 4 inputs: "
          f"{'OK' if ok else 'FAIL'}")
    return ok


def and3(a: int, b: int, c: int) -> int:
    """Two composed collisions. A at (2,2) and B at (7,7) collide into an A&B
    particle, which then collides with C launched at (0,9); the output cell
    (3,5) at step 4 is occupied iff a, b and c. Composing gates this way builds
    arbitrary circuits from the interaction gate."""
    H = W = 16
    g = [[0] * W for _ in range(H)]
    if a:
        g[2][2] = 1
    if b:
        g[7][7] = 1
    if c:
        g[0][9] = 1
    return run(g, 4)[3][5]


def _test_compose():
    ok = all(and3(a, b, c) == (a & b & c)
             for a in (0, 1) for b in (0, 1) for c in (0, 1))
    print(f"  composed 3-input AND (two chained collisions) over all 8 inputs: "
          f"{'OK' if ok else 'FAIL'}")
    return ok


def _test_collision():
    # Two particles interact (the joint evolution differs from independent
    # motion) and the collision stays reversible: the physics that logic needs.
    H = W = 12
    interacted = False
    revok = True
    # converging pairs: an SE-mover (even,even) meets an NW-mover (odd,odd) on a
    # shared diagonal, forming the 1001/0110 diagonal pair the rule deflects.
    for a, b in [((2, 2), (7, 7)), ((3, 3), (8, 8)), ((2, 8), (7, 3)),
                 ((4, 4), (9, 9)), ((2, 2), (9, 9))]:
        g = [[0] * W for _ in range(H)]
        g[a[0]][a[1]] = 1
        g[b[0]][b[1]] = 1
        ga = [[0] * W for _ in range(H)]
        ga[a[0]][a[1]] = 1
        gb = [[0] * W for _ in range(H)]
        gb[b[0]][b[1]] = 1
        joint = g
        for n in range(12):
            joint = step(joint, n & 1)
            ga = step(ga, n & 1)
            gb = step(gb, n & 1)
            free = _ball_positions(ga) | _ball_positions(gb)
            if _ball_positions(joint) != free:
                interacted = True
        if run_back(run(g, 12), 12) != g:
            revok = False
    print(f"  two-particle collisions interact (joint != independent motion): "
          f"{'OK' if interacted else 'FAIL'}")
    print(f"  collisions remain reversible: {'OK' if revok else 'FAIL'}")
    return interacted and revok


if __name__ == "__main__":
    print("Reversible Margolus cellular automaton")
    a = _test_rule()
    g = _test_gates()
    b = _test_reversibility()
    c = _test_ballistic()
    d = _test_collision()
    e = _test_gate()
    f = _test_compose()
    print("PASS" if (a and g and b and c and d and e and f) else "FAIL")