"""neural_computer8  --  the Collatz map iterated on the threshold CPU.

n -> n/2 if even, 3n+1 if odd, counting steps to 1. The parity test is a
bitwise AND with 1 that leaves the flags untouched, so the loop branches on the
comparison that precedes it. The register file is architecturally 8-bit, so the
machine computes the Collatz trajectory modulo 256; for seeds whose true
trajectory never exceeds 255, that is the exact Collatz step count, checked here
against a native reference.

    python demos/neural_computer8_collatz.py
"""
import os, sys, time
HERE = os.path.dirname(os.path.abspath(__file__))
REPO = os.path.dirname(HERE)
sys.path.insert(0, os.path.join(REPO, "src"))
sys.path.insert(0, os.path.join(REPO, "tools"))
from safetensors import safe_open
from cpu_programs import Asm
from eval_all import GenericThresholdCPU, get_manifest


def build(n0):
    a = Asm(1024)
    a.load(0, "N"); a.load(1, "ONE"); a.xor(2, 2)   # R0=n, R1=1, R2=steps
    a.label("loop")
    a.cmp(0, 1); a.jz("done")            # n == 1 ?
    a.xor(3, 3); a.add(3, 0); a.and_(3, 1)   # R3 = n & 1  (flags preserved)
    a.cmp(3, 1); a.jz("odd")
    a.shr(0); a.jmp("step")              # even: n >>= 1
    a.label("odd")
    a.xor(3, 3); a.add(3, 0); a.shl(3)   # R3 = 2n
    a.add(0, 3); a.add(0, 1)             # n = 3n + 1
    a.label("step")
    a.add(2, 1); a.jmp("loop")           # steps++
    a.label("done")
    a.store(2, "RES"); a.halt()
    a.org(0x300); a.label("N"); a.db(n0)
    a.label("ONE"); a.db(1); a.label("RES"); a.db(0)
    return a, a.assemble()


def true_collatz(n):
    peak, steps = n, 0
    while n != 1:
        n = n // 2 if n % 2 == 0 else 3 * n + 1
        peak = max(peak, n); steps += 1
    return steps, peak


if __name__ == "__main__":
    tens = {}
    with safe_open(os.path.join(REPO, "variants", "neural_computer8_small.safetensors"),
                   framework="pt") as f:
        for name in f.keys():
            tens[name] = f.get_tensor(name).float()
    cpu = GenericThresholdCPU(tens)
    addr = get_manifest(tens)["addr_bits"]

    print("neural_computer8: Collatz step counts through the gates")
    print("=" * 56)
    print("  seed   CPU steps   true steps   peak   in 8-bit range?")
    for n0 in (6, 7, 9, 18, 25, 45):
        a, mem = build(n0)
        state = {"pc": 0, "regs": [0] * 4, "flags": [0] * 4, "mem": list(mem),
                 "halted": False, "sp": (1 << addr) - 1}
        t0 = time.perf_counter()
        final, cycles = cpu.run(state, max_cycles=3000)
        steps = final["mem"][a.labels["RES"]]
        tsteps, peak = true_collatz(n0)
        inrange = peak <= 255
        match = "MATCH" if (inrange and steps == tsteps) else \
                ("mod-256" if not inrange else "MISMATCH")
        print(f"  {n0:4d}   {steps:9d}   {tsteps:10d}   {peak:4d}   "
              f"{'yes' if inrange else 'no (wraps)':10s} [{match}]")
    print("\n  For every seed whose trajectory stays under 256 the threshold CPU")
    print("  reproduces the exact Collatz step count.")