"""Run SUPERBEAM inside the constructed BeamForCausalLM (beamboot/chip8_lm.py) on the GPU.

  .venv-gpu/bin/python scripts/run_chip8_lm.py            # first 5 frames, verified
  .venv-gpu/bin/python scripts/run_chip8_lm.py --frames 171

Keys come from roms/superbeam_inputs.json (the keyboard device injects them as KEY tokens at
frame boundaries). Every token, every instruction's state (registers, stack, memory, display) and
every frame is checked against beamboot/chip8.py, and the frames also against
demo/superbeam_frames.json. Writes demo/superbeam_model_frames.json:
  {rom, ipf, frames_run, inputs, verified, mismatch, matches_superbeam_frames_json, n_tokens,
   prompt_tokens, decode_seconds, tok_per_s, device, model: {...}, frames: [[32 hex rows], ...]}
Prints tok/s and an ETA for all 171 frames (step time measured at longer cache lengths).
"""
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_lm as C  # noqa: E402

ROM = os.path.join(ROOT, 'roms', 'superbeam.ch8')
INPUTS = os.path.join(ROOT, 'roms', 'superbeam_inputs.json')
REF_FRAMES = os.path.join(ROOT, 'demo', 'superbeam_frames.json')
OUT = os.path.join(ROOT, 'demo', 'superbeam_model_frames.json')
FULL = 171


def total_tokens(rom, events, frames, ipf):
    """Stream length (prompt excluded) for `frames` frames, from the reference."""
    keys = C.key_tokens(events)
    n, nt = 0, 0
    for kind, ts, _ in C._ref_records(rom, ipf, 0, lambda fr: keys.get(fr, ())):
        n += len(ts)
        nt += kind == 'tick'
        if nt == frames:
            return n


def save(s, model, a, events, ref_frames):
    res = s.result(tokens=False)
    match = (res['frame_rows'] == ref_frames[:res['frames']]) if ref_frames else None
    out = dict(rom='roms/superbeam.ch8', inputs_file=os.path.relpath(a.inputs, ROOT), ipf=a.ipf, frames_run=res['frames'],
               inputs=[e for e in events if e[0] < res['frames']], verified=res['verified'],
               mismatch=res['mismatch'], matches_superbeam_frames_json=match, n_tokens=res['n_tokens'],
               prompt_tokens=res['prompt_tokens'], decode_seconds=round(res['decode_seconds'], 2),
               tok_per_s=round(res['tok_per_s'], 1), device=torch.cuda.get_device_name(),
               model=dict(cls=type(model).__name__, layers=model.config.num_hidden_layers,
                          hidden=model.config.hidden_size, vocab=model.config.vocab_size, dtype='float32',
                          attn=model.config._attn_implementation, decoder='GraphDecoder (CUDA graphs)'),
               frames=res['frame_rows'])
    with open(a.out + '.tmp', 'w') as f:
        json.dump(out, f)
    os.replace(a.out + '.tmp', a.out)
    print('wrote %s (%d frames)' % (os.path.relpath(a.out, ROOT), res['frames']), flush=True)
    res['matches'] = match
    return res


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--frames', type=int, default=5)
    ap.add_argument('--ipf', type=int, default=300)
    ap.add_argument('--no-eta', action='store_true')
    ap.add_argument('--inputs', default=INPUTS, help='key events list, or a schedule file with an "events" key')
    ap.add_argument('--out', default=OUT)
    ap.add_argument('--save-every', type=int, default=5, help='rewrite --out every N frames')
    a = ap.parse_args()
    rom = open(ROM, 'rb').read()
    events = json.load(open(a.inputs))
    if isinstance(events, dict):
        events = events['events']
    scripted = os.path.abspath(a.inputs) == os.path.abspath(INPUTS)
    ref_frames = json.load(open(REF_FRAMES))['frames'] if scripted else None

    t0 = time.perf_counter()
    model = C.get_model(a.ipf, 'cuda')
    print('model built in %.1fs: %s, %d layers, hidden %d, vocab %d, %.1fM params, %s'
          % (time.perf_counter() - t0, type(model).__name__, model.config.num_hidden_layers,
             model.config.hidden_size, model.config.vocab_size,
             sum(p.numel() for p in model.parameters()) / 1e6, torch.cuda.get_device_name()))
    s = C.Session(rom, 'cuda', a.ipf, verify=True, model=model)
    print('prompt %d tokens, prefill %.2fs' % (s.prompt_tokens, s.prefill_seconds))
    by = {}
    for g, k, act in events:
        by.setdefault(g, []).append((k, act == 'down'))
    for g in range(a.frames):
        n0, t1 = len(s.tokens), time.perf_counter()
        s.advance(by.get(g, ()))
        dt = time.perf_counter() - t1
        ok = 'ok' if s.mismatch is None else 'MISMATCH %s' % s.mismatch
        same = (s.result(tokens=False)['frame_rows'][-1] == ref_frames[g]) if ref_frames else 'n/a'
        print('frame %3d: %5d tokens %6.1fs %6.1f tok/s  verify %s  frame==superbeam_frames.json %s'
              % (g, len(s.tokens) - n0, dt, (len(s.tokens) - n0) / dt, ok, same))
        if s.mismatch is not None:
            break
        if a.save_every and (g + 1) % a.save_every == 0:
            save(s, model, a, events, ref_frames)
    res = save(s, model, a, events, ref_frames)
    match = res['matches']
    print('frames %d, tokens %d, decode %.1fs, %.1f tok/s, verified %s, frames match json %s'
          % (res['frames'], res['n_tokens'], res['decode_seconds'], res['tok_per_s'], res['verified'], match))

    if a.no_eta:
        return
    total = total_tokens(rom, events, FULL, a.ipf) + s.prompt_tokens
    done = s.prompt_tokens + len(s.tokens)
    pts = [done + 64] + [x for x in (16384, 32768, 65536, 131072) if done + 64 < x < total] + [total]
    times = [(L, s.dec.time_at(L)) for L in pts]
    print('step time vs cache length: ' + ', '.join('%d: %.1f ms (%.0f tok/s)' % (L, 1e3 * t, 1 / t)
                                                    for L, t in times))
    eta = 0.0
    for (l0, t0_), (l1, t1_) in zip(times, times[1:]):
        eta += (l1 - l0) * (t0_ + t1_) / 2
    print('171 frames = %d tokens (incl. prompt); remaining ETA %.1f min, whole run ~%.1f min'
          % (total, eta / 60, (eta + res['decode_seconds']) / 60))


if __name__ == '__main__':
    main()