#!/usr/bin/env python3
"""Replay a SUPERBEAM key schedule inside the constructed BeamForCausalLM, fast path
(beamboot/chip8_fast.py: per-frame checkpoints + speculative verification).

  .venv-gpu/bin/python scripts/run_chip8_fast.py --inputs results/beam_key_schedule.json \
      --frames 175 --out demo/beam_model_frames.json

Writes the same JSON shape as scripts/run_chip8_lm.py (frames = 32 hex rows per frame) plus
per-frame stats, and compares every frame with the reference emulator (superbeam_gym replay).
"""
import argparse
import json
import os
import sys
import time

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

import torch  # noqa: E402

from beamboot import chip8_fast as F  # noqa: E402
from beamboot.superbeam_gym import SuperBeamEnv, frames_from_events  # noqa: E402


def emulator_rows(events, n):
    env = SuperBeamEnv(backend='emulator')
    env.reset()
    out = []
    for keys in frames_from_events(events, n):
        env.backend.advance(keys)
        out.append(env.backend.m.display_rows())
    return out


def save(path, s, ref, args, wall):
    st = s.stats
    k = 0
    while k < len(s.frame_rows) and s.frame_rows[k] == ref[k]:
        k += 1
    secs = [x.seconds for x in st]
    out = dict(rom='roms/superbeam.ch8', inputs_file=os.path.relpath(args.inputs, ROOT), ipf=s.ipf,
               frames_run=len(s.frame_rows), verified=k == len(s.frame_rows),
               frames_equal_emulator=k, mismatch=None if k == len(s.frame_rows) else k,
               method='per-frame checkpoint + speculative verification (beamboot/chip8_fast.py)',
               frames_all_draft_accepted=sum(x.accepted_all for x in st),
               n_tokens=sum(x.frame_tokens for x in st), prompt_tokens_per_frame=st[-1].prompt_tokens,
               seconds_per_frame=round(sum(secs) / len(secs), 3), max_seconds_per_frame=round(max(secs), 3),
               wall_seconds=round(wall, 1),
               tok_per_s=round(sum(x.frame_tokens for x in st) / sum(secs), 1),
               device=torch.cuda.get_device_name() if torch.cuda.is_available() else 'cpu',
               model=dict(cls=type(s.model).__name__, layers=s.model.config.num_hidden_layers,
                          hidden=s.model.config.hidden_size, vocab=s.model.config.vocab_size,
                          dtype='float32', attn='sdpa'),
               per_frame=[dict(f=x.frame, tokens=x.frame_tokens, accepted_all=x.accepted_all,
                               first_reject=x.first_reject, s=round(x.seconds, 3)) for x in st],
               frames=s.frame_rows)
    with open(path + '.tmp', 'w') as f:
        json.dump(out, f)
    os.replace(path + '.tmp', path)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--inputs', default=os.path.join(ROOT, 'results', 'beam_key_schedule.json'))
    ap.add_argument('--frames', type=int, default=None)
    ap.add_argument('--out', default=os.path.join(ROOT, 'demo', 'beam_model_frames.json'))
    ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
    args = ap.parse_args()
    sched = json.load(open(args.inputs))
    events = sched['events'] if isinstance(sched, dict) else sched
    n = args.frames or (sched.get('frames') if isinstance(sched, dict) else None)
    if not n:                                   # plain event list: run until the emulator's win
        n = SuperBeamEnv(backend='emulator', frame_skip=1).replay_schedule(events)['frames']
    rom = open(os.path.join(ROOT, 'roms', 'superbeam.ch8'), 'rb').read()
    ref = emulator_rows(events, n)
    t0 = time.time()

    def on_frame(g, s):
        x = s.stats[-1]
        ok = s.frame_rows[-1] == ref[g]
        print('frame %3d: %4d tokens  %.2fs  draft %s  frame==emulator %s'
              % (g, x.frame_tokens, x.seconds, 'accepted' if x.accepted_all else 'REJECTED at %d' % x.first_reject,
                 ok if ok else 'MISMATCH'), flush=True)
        if (g + 1) % 10 == 0:
            save(args.out, s, ref, args, time.time() - t0)

    s = F.run(rom, events, frames=n, device=args.device, on_frame=on_frame)
    out = save(args.out, s, ref, args, time.time() - t0)
    print('done: %d frames, equal to emulator %d, all drafts accepted %d, %.2f s/frame, %.0f tok/s, wall %.0fs'
          % (out['frames_run'], out['frames_equal_emulator'], out['frames_all_draft_accepted'],
             out['seconds_per_frame'], out['tok_per_s'], out['wall_seconds']))
    return 0 if out['verified'] else 1


if __name__ == '__main__':
    raise SystemExit(main())