import os

os.environ.setdefault("CHIPFORMER_WROM_HEADS", "320")
import sys
import time

sys.path.insert(0, "/tmp/opencode/qp/experimental/riscv")
from chipformer.rv import elf, ref

m = ref.Machine.from_program(elf.program(sys.argv[1]))
n = int(sys.argv[2])
t0 = time.time()
ntok = 0
while not m.halted:
    ntok += len(m.record())
d = lambda i: ref.sx(m.ram[0xE000 // 4 + i], 32)
print("instrucciones", m.steps, "trace tokens", ntok, f"{time.time() - t0:.0f}s")
print("kernel tokens :", [d(100 + p) for p in range(n)])
print("kernel logits :", [d(200 + p) for p in range(n)])