"""The reference machine against Definition 3.1, and the halt fixed point.

Definition 3.1 specifies one SUBLEQ step: read A = M[p], B = M[p+1], C = M[p+2],
set M[B] <- M[B] - M[A] mod 256, and set p <- C when the new M[B] is at most 0
in two's complement and p <- p + 3 otherwise, with indices taken modulo 256 and
a halt at p = 255. The clauses are transcribed here independently of
selfrep.run_reference and the two are compared, with the program counter swept
over every value so the wrapped fetches at the top of memory are covered.

Proposition 3.10 asserts in addition that a halted state is a fixed point of the
transition, because every architectural write is gated by the negation of the
halt bit. That is checked on the levelized map itself, not on the reference.
"""
import json
import os
import random
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from selfrep import Tape, HALT_PC

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def runs_path(name: str) -> str:
    d = os.path.join(REPO, "paper", "runs")
    os.makedirs(d, exist_ok=True)
    return os.path.join(d, name)


def spec_step(mem, p):
    """Definition 3.1, transcribed."""
    A = mem[p % 256]
    B = mem[(p + 1) % 256]
    C = mem[(p + 2) % 256]
    m = list(mem)
    r = (m[B % 256] - m[A % 256]) % 256
    m[B % 256] = r
    two_c = r - 256 if r >= 128 else r
    return m, (C if two_c <= 0 else (p + 3) % 256)


def impl_step(mem, p):
    """The reference of selfrep.run_reference, one step, device detached."""
    m = list(mem)
    A = m[p]
    B = m[(p + 1) & 0xFF]
    C = m[(p + 2) & 0xFF]
    r = (m[B] - m[A]) & 0xFF
    m[B] = r
    return m, (C if (r == 0 or r >= 0x80) else (p + 3) & 0xFF)


def main() -> int:
    rng = random.Random(31)
    bad, n = [], 0
    mems = [[0] * 256, [255] * 256, list(range(256)), list(reversed(range(256)))]
    mems += [[rng.randrange(256) for _ in range(256)] for _ in range(60)]
    for mem in mems:
        for p in range(256):                       # every counter, wraps included
            if p == HALT_PC:
                continue
            sm, sp = spec_step(mem, p)
            im, ip = impl_step(mem, p)
            n += 1
            if sm != im or sp != ip:
                bad.append((p, sp, ip))
    print(f"  reference steps compared with Definition 3.1: {n:,}")
    print(f"  agree on every one: {'yes' if not bad else f'NO {bad[:4]}'}")

    wrapped = sum(1 for p in (254, 255, 253) if p != HALT_PC) * len(mems)
    print(f"  of which {len(mems) * 3:,} fetch an operand across the top of memory")

    # A halted state is a fixed point of the levelized map. The state is halted
    # when the halt bit is set, coordinate 8 of the vector; the bit latches one
    # step after the next counter decodes to 255, so setting the counter alone
    # does not halt the machine.
    from selfrep import LevEvaluator, read_host
    import torch
    L = LevEvaluator(read_host(), device="cpu", dense=False)
    gen = torch.Generator().manual_seed(97)
    V = (torch.rand(48, L.N, generator=gen) < 0.5).float()
    V[:, 8] = 1.0
    fixed = bool((L.step(V) == V).all())
    print(f"  halted states are fixed points of Lev(N_host), 48 random states: "
          f"{'yes' if fixed else 'NO'}")

    # And the bit latches exactly when the next counter is 255: run one step
    # from unhalted states and compare the latched bit against the reference.
    U = (torch.rand(64, L.N, generator=gen) < 0.5).float()
    U[:, 8] = 0.0
    W = L.step(U)
    latch_bad = 0
    for i in range(U.shape[0]):
        mem = [int(sum(int(U[i, 9 + j * 8 + k]) << (7 - k) for k in range(8)))
               for j in range(256)]
        pc = int(sum(int(U[i, k]) << (7 - k) for k in range(8)))
        _, npc = impl_step(mem, pc)
        if int(W[i, 8]) != (1 if npc == HALT_PC else 0):
            latch_bad += 1
    print(f"  the halt bit latches exactly when the next counter is 255, "
          f"64 random states: {'yes' if latch_bad == 0 else f'NO ({latch_bad})'}")
    fixed = fixed and latch_bad == 0

    ok = not bad and fixed
    json.dump({"steps_compared": n, "disagreements": len(bad),
               "halt_fixed_point": fixed},
              open(runs_path("paper_host.json"), "w"), indent=1)
    return 0 if ok else 1


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