"""neural_matrix8  --  a fleet of processors as one batched matrix product.

The whole CPU is compiled to 108 ternary weight matrices with a Heaviside step
between them, so one clock cycle is one matrix-vector product plus a threshold.
Add a batch dimension and N processors are the same 108 matmuls: here 65,536
independent CPUs step in lockstep on the GPU. A 256-wide correctness fleet
confirms every CPU halts exactly at cycle 2n+1 for its own countdown input.

    python demos/neural_matrix8_gpu_cpu_fleet.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 matrix8 import MatrixMachine, state_to_vec, _mk_state, _instr


def prog_bytes():
    return (_instr(0x1, 0, 1)        # SUB R0, R1
            + _instr(0xD, 0, 0, 1) + [0, 0]   # JNZ 0x0000 (address word)
            + _instr(0xF)            # HALT
            + [0] * 8)


if __name__ == "__main__":
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    mm = MatrixMachine.from_file(device=dev)
    n_neurons = sum(int(b.numel()) for b in mm.B)
    name = torch.cuda.get_device_name(0) if dev == "cuda" else "cpu"
    print("neural_matrix8: 65,536 CPUs as one batched matrix product")
    print("=" * 56)
    print(f"loaded {len(mm.W)} ternary matrices, {n_neurons:,} threshold neurons "
          f"per transition, device={dev} ({name})")
    mem = prog_bytes()

    # correctness fleet: one CPU per 8-bit input, all 256 at once
    V = torch.stack([state_to_vec(_mk_state(mem=mem, regs=(x, 1, 0, 0)))
                     for x in range(256)]).to(dev)
    first_halt = torch.full((256,), -1, dtype=torch.long)
    step = 0
    while step < 600:
        halted = V[:, MatrixMachine.HALT_IDX] > 0.5
        first_halt[(first_halt < 0) & halted.cpu()] = step
        if bool(halted.all()):
            break
        V = mm.step(V); step += 1
    expect = torch.tensor([2 * (x if x else 256) + 1 for x in range(256)])
    ok_halt = bool((first_halt == expect).all())
    ok_r0 = bool((V[:, 4:12].cpu().long().sum(1) == 0).all())
    print(f"\ncorrectness fleet of 256: every CPU halts at cycle 2n+1 for its own "
          f"input: {'EXACT' if ok_halt else 'MISMATCH'}; all results R0==0: {ok_r0}")

    # throughput fleet: 65,536 CPUs
    B = 65536
    V = torch.stack([state_to_vec(_mk_state(mem=mem, regs=(x & 0xFF, 1, 0, 0)))
                     for x in range(256)]).repeat(256, 1).to(dev)
    if dev == "cuda":
        torch.cuda.synchronize()
    t0, steps = time.perf_counter(), 0
    while steps < 520:
        V = mm.step(V); steps += 1
        if steps % 64 == 0 and bool((V[:, MatrixMachine.HALT_IDX] > 0.5).all()):
            break
    if dev == "cuda":
        torch.cuda.synchronize()
    dt = time.perf_counter() - t0
    instr = B * steps
    print(f"\nthroughput fleet of {B:,}: {steps} transitions in {dt:.1f}s")
    print(f"  {instr / dt / 1e6:,.1f} million instructions/s aggregate "
          f"({instr:,} instructions retired)")
    print(f"  {instr * n_neurons / dt / 1e9:,.1f} billion threshold-neuron "
          f"evaluations/s")