"""neural_reflect  --  three levels of machine in one tensor.

Level 0 is U, a fixed ternary threshold interpreter (~24k gates). Level 1 is a
complete SUBLEQ computer stored as data inside U's writable state. Level 2 is a
SUBLEQ program that rewrites its own operand bytes as it runs: a zeroing
instruction that walks itself across memory, erasing a block. The interpreter,
the machine it runs, and the program that machine executes are all the same
tensor; the final memory is checked bit-for-bit against a native SUBLEQ emulator.

    python demos/neural_reflect_self_modifying_stack.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"))
import torch
from reflect import (Cfg, build_subleq_machine, build_net, Leveled,
                     encode_netlist, state_to_vec, subleq32_step, MEM, MEMB)


if __name__ == "__main__":
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    cfg = Cfg(A=14, G=296, banks=1, self_mod=False)
    prog, lay = build_subleq_machine(cfg)
    inet, ii, io = build_net(cfg)
    print("neural_reflect: self-modifying code on a stored machine on U")
    print("=" * 60)
    print(f"level 0: interpreter U = {len(inet.gates):,} threshold gates (fixed)")
    print(f"level 1: a SUBLEQ machine = {len(prog)} gates, stored as writable data")
    lev = Leveled(inet, ii, io, device=dev)
    nl = encode_netlist(cfg, prog)
    sig0 = [0] * cfg.S
    sig0[cfg.NET0:cfg.NET0 + len(nl)] = nl
    base = state_to_vec(cfg, {"sig": sig0, "gp": 0, "halt": 0}).to(dev)
    PC = lay["PC"]

    # level 2: SUBLEQ (A B C: M[B]-=M[A]; if <=0 goto C). The first instruction
    # zeroes a target cell, then two instructions increment its own A and B
    # operand bytes, so the zeroing walks across a data block.
    pm = [0] * 32
    pm[0:3]   = [24, 24, 3]      # zero M[<target>]; these operands get rewritten
    pm[3:6]   = [19, 0, 6]       # M[0] += 1  (patch own A-field: cell 0)
    pm[6:9]   = [19, 1, 9]       # M[1] += 1  (patch own B-field: cell 1)
    pm[9:12]  = [20, 21, 15]     # count -= 1; if <= 0 branch to epilogue
    pm[12:15] = [18, 18, 0]      # 0 -> unconditional jump to loop head
    pm[15:18] = [18, 18, 31]     # epilogue: branch to 31 = HALT
    pm[18], pm[19], pm[20], pm[21] = 0, 255, 1, 4
    pm[24:28] = [7, 99, 123, 200]          # the block the program will erase
    print(f"level 2: a program whose instruction 0 starts as [24,24,3] and "
          f"rewrites its own operands")

    ref_mem, ref_pc = list(pm), 0
    while ref_pc != 31:
        ref_mem, ref_pc = subleq32_step(ref_mem, ref_pc)

    V = base.unsqueeze(0).clone()
    for i in range(32):
        for b in range(8):
            V[:, MEM + b * MEMB + i] = float((pm[i] >> b) & 1)

    def get_mem(Vc):
        return sum(Vc[:, MEM + b * MEMB: MEM + b * MEMB + 32].long() << b
                   for b in range(8))[0].tolist()

    t0, walk, steps, halted = time.perf_counter(), [], 0, False
    for _ in range(26):
        m_now = get_mem(V.cpu())
        walk.append((m_now[0], m_now[1]))
        for _ in range(cfg.G):
            V = lev.step(V)
        steps += 1
        if int(V[0, cfg.S + cfg.GPW]) == 1:
            halted = True
            break
    dt = time.perf_counter() - t0
    got = get_mem(V.cpu())

    mut = " -> ".join(f"[{a},{b}]" for a, b in
                      [walk[0]] + [w for i, w in enumerate(walk[1:], 1) if w != walk[i - 1]])
    print(f"\nran {steps} hosted instructions in {dt:.0f}s "
          f"({steps * cfg.G} interpreter recurrences); halted={halted}")
    print(f"  instruction 0's operand cells as it ran: {mut}")
    print(f"  target block M[24..27]: {pm[24:28]} -> {got[24:28]}")
    print(f"  final memory == native SUBLEQ reference: "
          f"{'EXACT' if got == ref_mem else 'MISMATCH'}")
    print("\n  The program, the machine it runs on, and the machine that runs")
    print("  that machine are all the same tensor.")