"""Probe the BeamAttention RAM (beamboot/attn_ram.py) at long horizons.

Seeded random writes/reads vs a dict. Clock = row position, T = 2^18 (or the
next power of two >= horizon). Reports read errors, the minimum logit margin
(true winner minus best other key, fp32 recompute from the real cache keys and
the module's own q_norm/q_proj/query scaling), rows/s of the attention forward
alone, and peak GPU memory.

  .venv-gpu/bin/python scripts/probe_attn_ram.py --horizons 4096 32768 262144
"""
import argparse
import json
import random
import sys
import time
from pathlib import Path

import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from beamboot.attn_ram import AttnRAM, HEAD  # noqa: E402
from transformers.models.beam.modeling_beam import compute_attention_scaling  # noqa: E402


def make_stream(rng, n):
    """Writes ~70% (bursts to one address model CHIP-8 V/I/stack traffic), reads ~30%."""
    ev, hot = [], [rng.randrange(4096) for _ in range(96)]
    while len(ev) < n:
        r = rng.random()
        if r < 0.10:                                  # burst: same address written on adjacent clocks
            a = rng.choice(hot)
            for _ in range(rng.randrange(2, 5)): ev.append(('w', a, rng.randrange(256)))
            ev.append(('r', a))
        elif r < 0.70:
            a = rng.choice(hot) if rng.random() < 0.6 else rng.randrange(4096)
            ev.append(('w', a, rng.randrange(256)))
        else:
            a = rng.choice(hot) if rng.random() < 0.6 else rng.randrange(4096)
            ev.append(('r', a))
    return ev[:n]


def probe(horizon, impl, dtype, device, chunk, seed, margins=True):
    T = max(2 ** 18, 1 << (horizon - 1).bit_length())
    torch.cuda.reset_peak_memory_stats() if device == 'cuda' else None
    ram = AttnRAM(max_clock=T, device=device, dtype=dtype, attn_impl=impl)
    events = make_stream(random.Random(seed), horizon - 1)
    latest, ref = {}, {}                   # addr -> position of latest write / value
    errors = reads = 0
    min_margin, fwd_time = float('inf'), 0.0
    L = ram.layer
    for start in range(0, len(events), chunk):
        block = events[start:start + chunk]
        past = ram.length
        win, want, qpos, qaddr = [], [], [], []
        for i, e in enumerate(block):
            p = past + i
            if e[0] == 'w': latest[e[1]] = p; ref[e[1]] = e[2]
            else:
                win.append(latest.get(e[1], 0)); want.append(ref.get(e[1], 0))
                qpos.append(p); qaddr.append(e[1])
        if device == 'cuda': torch.cuda.synchronize()
        t0 = time.perf_counter()
        got = ram.run(block)
        if device == 'cuda': torch.cuda.synchronize()
        fwd_time += time.perf_counter() - t0
        reads += len(got)
        errors += sum(g != w for g, w in zip(got, want))
        if margins and qpos:
            x = torch.stack([ram.read_row(a) for a in qaddr]).to(device, dtype)[None]
            with torch.inference_mode():
                q = L.q_norm(L.q_proj(x).view(1, -1, 1, HEAD)).transpose(1, 2).float()
                pos = torch.tensor([qpos], device=device)
                q = (q * compute_attention_scaling(pos))[0, 0]
                k = ram.cache.layers[0].keys[0, 0].float()
                s = (q @ k.T) * L.scaling
                cols = torch.arange(k.shape[0], device=device)[None]
                s.masked_fill_(cols > pos[0][:, None], float('-inf'))
                w = torch.tensor(win, device=device)
                ws = s.gather(1, w[:, None])[:, 0]
                s.scatter_(1, w[:, None], float('-inf'))
                min_margin = min(min_margin, float((ws - s.max(1).values).min()))
    res = dict(horizon=horizon, max_clock=T, impl=impl, dtype=str(dtype).split('.')[-1], device=device,
               rows=ram.length, reads=reads, errors=errors,
               min_margin=round(min_margin, 2) if margins else None,
               rows_per_s=round(ram.length / fwd_time, 1), fwd_seconds=round(fwd_time, 2))
    if device == 'cuda': res['peak_gpu_MiB'] = round(torch.cuda.max_memory_allocated() / 2 ** 20)
    del ram
    if device == 'cuda': torch.cuda.empty_cache()
    return res


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--horizons', type=int, nargs='+', default=[4096, 32768, 262144])
    ap.add_argument('--impls', nargs='+', default=['sdpa', 'eager'])
    ap.add_argument('--dtypes', nargs='+', default=['float32', 'bfloat16'])
    ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
    ap.add_argument('--chunk', type=int, default=512)
    ap.add_argument('--seed', type=int, default=18)
    a = ap.parse_args()
    for dt in a.dtypes:
        for h in a.horizons:
            for impl in a.impls:
                try:
                    r = probe(h, impl, getattr(torch, dt), a.device, a.chunk, a.seed)
                except torch.OutOfMemoryError as exc:
                    r = dict(horizon=h, impl=impl, dtype=dt, error=f'OOM: {str(exc)[:80]}')
                print(json.dumps(r), flush=True)


if __name__ == '__main__':
    main()