"""Find concrete RAM reads in the in-weights CHIP-8 computer (beamboot/chip8_lm.py).

Runs the constructed BeamForCausalLM on one real SUPERBEAM frame (checkpoint prompt + the frame's
tokens, as chip8_fast does) with eager attention. For the RAM-read layers (L4: opcode and memory
reads, L5: Vx/Vy reads) it lists query positions whose attention sits on a single write token while
older writes to the same address exist earlier in the context: the "latest-write RAM" circuit,
read straight out of the model's attention weights.

  .venv-gpu/bin/python scripts/attention_probe.py --frames 6 --out demo/attention_probe.json
"""
import argparse
import json
import os
import sys

import torch

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)

from beamboot import chip8_fast as F  # noqa: E402
from beamboot import chip8_lm as C    # noqa: E402


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--frames', type=int, default=6, help='advance this many frames first')
    ap.add_argument('--draft-tokens', type=int, default=600)
    ap.add_argument('--out', default=os.path.join(ROOT, 'demo', 'attention_probe.json'))
    args = ap.parse_args()

    rom = open(os.path.join(ROOT, 'roms', 'superbeam.ch8'), 'rb').read()
    sched = json.load(open(os.path.join(ROOT, 'roms', 'superbeam_inputs.json')))
    s = F.run(rom, [e for e in sched if int(e[0]) < args.frames], frames=args.frames, device='cpu')
    for g, k, a in sched:
        if int(g) == args.frames:
            s.c.set_key(int(k) & 15, a == 'down')
    p = F.checkpoint_prompt(s.c, s.ipf)
    d = s._draft([])[:args.draft_tokens]
    seq = p + d
    m = s.model
    m.config._attn_implementation = 'eager'
    with torch.inference_mode():
        out = m(input_ids=torch.tensor([seq]), output_attentions=True, logits_to_keep=len(d) + 1)
    am = out.logits[0].argmax(-1).tolist()
    agree = sum(1 for j, t in enumerate(d) if am[j] == t)
    labels = [tuple(C.untok(t)) for t in seq]

    def addr(lb):
        return lb[0], (lb[2] if lb[0] == 'VW' and len(lb) > 2 else None)

    mem = s.c.mem                                    # code bytes do not change during the frame

    def pc_at(q):                                    # PC of the record a query belongs to
        return next(int(labels[i][1]) for i in range(q - 1, -1, -1) if labels[i][0] == 'PCW')

    found = []
    for li in (4,):                                  # L5: Vx / Vy reads
        A = out.attentions[li][0]                    # (heads, q, k)
        for h in range(A.shape[0]):
            for q in range(len(p), len(seq)):
                if labels[q][0] != 'CLKL':           # the record's plan is computed at its CLKL token
                    continue
                w = A[h, q, :q + 1]
                kpos = int(w.argmax())
                if float(w[kpos]) < 0.99 or labels[kpos][0] not in ('VW', 'PCW', 'IW'):
                    continue
                if any(addr(labels[i]) == addr(labels[kpos]) for i in range(kpos + 1, q)):
                    continue                         # not the latest write before the query
                same = [i for i in range(kpos) if addr(labels[i]) == addr(labels[kpos])]
                if len(same) < 2:
                    continue
                pc = pc_at(q)
                found.append(dict(layer=li + 1, head=h, query=q, pc=pc, opcode='%02X%02X' % (mem[pc], mem[pc + 1]),
                                  older_vals=[str(labels[i][1]) for i in same[-6:]], query_tok=[str(x) for x in labels[q]],
                                  key=kpos, key_tok=[str(x) for x in labels[kpos]], weight=float(w[kpos]),
                                  older_same=len(same), older_pos=same[-20:],
                                  older_weight=float(w[same].sum()) if same else 0.0, ctx=q + 1))
    found.sort(key=lambda f: (-f['older_same'], f['layer']))

    # Head 0 of L5 is the Vx read. Check it on every instruction of the frame that uses Vx (x = second
    # nibble): its argmax must be the newest VW write to register x before the query (or, if x was
    # never written in context, no VW of x at all).
    A0 = out.attentions[4][0][0]
    uses_vx = set('345678ABEF') - {'A', 'B'}
    checked = correct = 0
    for q in range(len(p), len(seq)):
        if labels[q][0] != 'CLKL':
            continue
        pc = pc_at(q)
        op = '%02X%02X' % (mem[pc], mem[pc + 1])
        if op[0] not in uses_vx:
            continue
        x = int(op[1], 16)
        writes = [i for i in range(q) if labels[i][0] == 'VW' and int(labels[i][2]) == x]
        if not writes:
            continue
        checked += 1
        kpos = int(A0[q, :q + 1].argmax())
        correct += int(kpos == writes[-1] and float(A0[q, kpos]) > 0.99)
    stats = dict(head='L5 head 0 (Vx read)', instructions_checked=checked, read_newest_write=correct)

    fig = None
    best = [f for f in found if f['head'] == 0 and f['opcode'][0] in '34' and f['older_same'] >= 5]
    if best:
        f = best[0]
        q, x = f['query'], int(f['key_tok'][2])
        w = A0[q, :q + 1]
        writes = [i for i in range(q) if labels[i][0] == 'VW' and int(labels[i][2]) == x]
        fig = dict(query=q, ctx=q + 1, pc=f['pc'], opcode=f['opcode'], reg=x,
                   writes=[dict(pos=i, val=int(labels[i][1]), w=float(w[i])) for i in writes],
                   other_mass=float(w.sum() - w[writes].sum()),
                   kinds=[labels[i][0] for i in range(q + 1)])
    res = dict(frame=args.frames, prompt_tokens=len(p), checked_tokens=len(d), draft_agree=agree,
               n_reads=len(found), vx_read=stats, figure=fig, examples=found[:40])
    print(json.dumps(stats))
    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    json.dump(res, open(args.out, 'w'), indent=1)
    print(json.dumps({k: v for k, v in res.items() if k != 'examples'}))
    for f in found[:10]:
        print({k: v for k, v in f.items() if k != 'older_pos'})


if __name__ == '__main__':
    main()