"""neural_computer8  --  the Sieve of Eratosthenes with self-modifying code.

The threshold CPU's ISA has no indexed addressing, so the program does exactly
what memory-constrained 1970s code did: it rewrites the address bytes of its own
LOAD and STORE instructions before each access, walking a pointer across the
flags array. Every gate that fetches, decodes, marks, and patches is a threshold
neuron. It halts with all 54 primes below 256 marked.

    python demos/neural_computer8_self_modifying_sieve.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, _enc
from eval_all import GenericThresholdCPU, get_manifest

FLAGS = 0x300                       # flags[0..255] at 0x300..0x3FF


def build():
    g = Asm(1024)
    g.load(3, "ONE")                    # R3 = 1 (mark value / increment)
    g.load(0, "TWO")                    # R0 = p = 2
    g.label("outer")
    g.store(0, "Rlo")                   # patch probe address low byte <- p
    g.dw(_enc(0xA, 2, 0))               # probe: R2 = flags[p]  (self-modified)
    g.db(FLAGS >> 8); g.label("Rlo"); g.db(0x00)
    g.cmp(2, 3); g.jz("next_p")         # flags[p] == 1 -> composite, skip
    g.xor(1, 1); g.add(1, 0); g.add(1, 0)   # m = 2p
    g.label("inner")
    g.store(1, "Wlo")                   # patch mark address low byte <- m
    g.dw(_enc(0xB, 0, 3))               # mark: flags[m] = R3 = 1  (self-modified)
    g.db(FLAGS >> 8); g.label("Wlo"); g.db(0x00)
    g.add(1, 0)                         # m += p  (carry iff we passed 255)
    g.jnc("inner")
    g.label("next_p")
    g.add(0, 3)                         # p += 1
    g.load(2, "SIXTEEN")
    g.cmp(0, 2)
    g.jnz("outer")                      # loop until p == 16 (16^2 > 255)
    g.halt()
    g.org(0x200)
    g.label("ONE"); g.db(1)
    g.label("TWO"); g.db(2)
    g.label("SIXTEEN"); g.db(16)
    return g, g.assemble()


if __name__ == "__main__":
    g, mem = build()
    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)
    state = {"pc": 0, "regs": [0] * 4, "flags": [0] * 4, "mem": list(mem),
             "halted": False, "sp": (1 << get_manifest(tens)["addr_bits"]) - 1}
    t0 = time.perf_counter()
    final, cycles = cpu.run(state, max_cycles=4000)
    dt = time.perf_counter() - t0

    got = [n for n in range(2, 256) if final["mem"][FLAGS + n] == 0]
    ref = [n for n in range(2, 256) if all(n % d for d in range(2, int(n**0.5) + 1))]
    print("neural_computer8: self-modifying Sieve of Eratosthenes")
    print("=" * 56)
    print(f"halted={final['halted']}  {cycles} cycles through the gates ({dt:.0f}s)")
    print(f"primes < 256 found: {len(got)}   native sieve: {len(ref)}   "
          f"{'EXACT MATCH' if got == ref else 'MISMATCH'}")
    print("  " + " ".join(map(str, got)))
    print(f"self-modified operand bytes at halt: probe={final['mem'][g.labels['Rlo']]}, "
          f"mark={final['mem'][g.labels['Wlo']]} (the program rewrote its own "
          f"instruction stream as it ran)")