import collections
import copy
import os
import subprocess
import sys
import time

os.environ.setdefault("CHIPFORMER_WROM_HEADS", "320")
sys.path.insert(0, "/tmp/opencode/qp/experimental/riscv")
from chipformer.q35 import ir
from chipformer.rv import core, elf, ref
from chipformer.rv import spec as S

path, where, ntok = sys.argv[1], sys.argv[2], int(sys.argv[3])
syms = {}
for line in subprocess.run(["llvm-nm", path], capture_output=True, text=True).stdout.splitlines():
    parts = line.split()
    if len(parts) == 3:
        syms[parts[2]] = int(parts[0], 16)
m = ref.Machine.from_program(elf.program(path))
dump = lambda i: m.ram[0xE000 // 4 + i]


def op_at(mm):
    if mm.n == S.N_MAX:
        return None
    return ref.decode(mm.rom[(mm.pc - S.ROM_BASE) >> 2]).op


def advance(pred):
    while not m.halted and not pred():
        m.record()
    assert not m.halted, "el programa termino antes del punto de control"


if where == "start":
    pass
elif where == "bank":
    advance(lambda: m.bank >= 1)
elif where == "down":
    advance(lambda: m.pc == syms["vk_8_init_raw"])
elif where == "pos1_attn":
    advance(lambda: dump(100) != 0)
    advance(lambda: op_at(m) == "SMAQA")
elif where == "head_bank":
    advance(lambda: dump(30) != 0)
    b0 = m.bank
    advance(lambda: m.bank != b0)
print(f"punto de control {where}: instrucciones ejecutadas {m.steps}, pc {m.pc:#x}, banco {m.bank}, wb {m.wb}", flush=True)
ref_m = copy.deepcopy(m)
mix = collections.Counter()
toks = [S.tok("GO")]
while len(toks) < ntok and not ref_m.halted:
    mix[op_at(ref_m)] += 1
    toks += ref_m.record()
print("mezcla de instrucciones del tramo:", dict(mix.most_common(14)), flush=True)
t0 = time.time()
c = core.build()
core.snapshot(c, m)
mat = ir.evaluate(c, toks, check=True)
pred = core.predictions(c, mat)
bad = [(i, S.name(toks[i + 1]), pred[i]) for i in range(len(toks) - 1) if pred[i] != toks[i + 1]]
fams = collections.Counter(S.family(t)[0] for t in toks)
print(f"tokens {len(toks)}, tokens WB {fams['WB']}, WBB {fams['WBB']}, desacuerdos {len(bad)}, {time.time() - t0:.0f}s", flush=True)
print("CIRCUITO == REFERENCIA" if not bad else f"FALLA: {bad[:5]}", flush=True)