"""MoBA benchmark: 20M context, target 50k+ tok/s.

Tests MoBA (Mixture of Block Attention) at 20M token context.
Block-sparse: each query attends to top-k blocks instead of all 20M tokens.

Physics:
  Dense attention at 20M: ~4,656 tok/s (reads 10.2 GB KV per step)
  MoBA top-k=3, block_size=4096: reads 3*4096*512B = 6 MB per step (L2 cache!)
  Theoretical: millions of tok/s (compute-bound, not memory-bound)
"""
import time
import torch
import sys

from config import Config
from model import build_model

print(f"GPU: {torch.cuda.get_device_name(0)}")
vram_total = torch.cuda.mem_get_info()[1] / 1e9
vram_free = torch.cuda.mem_get_info()[0] / 1e9
print(f"VRAM: {vram_free:.1f} GB free / {vram_total:.1f} GB total")

# --- Config: MoBA at 20M context ---
CTX = 20_000_000
K = 4096
K1 = K + 1
BLOCK_SIZE = K1  # = K+1 = 4097: each step fills exactly one block
TOP_K = 3
N_BLOCKS = (CTX // BLOCK_SIZE) + 2  # +1 for current partial, +1 safety

cfg = Config(
    d_model=384,
    n_layers=1,
    n_heads=12,
    n_kv_heads=4,
    ffn_mult=2.667,
    medusa_heads=K,
    medusa_rank=1,
    medusa_hidden=0,
    window_size=0,           # disable sliding window — use MoBA instead
    max_seq_len=CTX + K + 100,
    use_alibi=False,         # no position encoding needed for MoBA
    moba_block_size=BLOCK_SIZE,
    moba_top_k=TOP_K,
    dtype="bfloat16",
)

dt = getattr(torch, cfg.dtype)
dev = "cuda"

# --- VRAM estimate ---
kv_per_token = 2 * cfg.n_kv_heads * cfg.head_dim * 2 * 2  # K+V, bf16, per layer
kv_gb = CTX * kv_per_token / 1e9
kbar_gb = N_BLOCKS * cfg.n_kv_heads * cfg.head_dim * 2 / 1e9
params = cfg.count_params()
params_gb = params["grand_total"] * 2 / 1e9
est_total = kv_gb + kbar_gb + params_gb + 1.0  # +1 GB headroom for activations

print(f"\nConfig: d={cfg.d_model} L={cfg.n_layers} H={cfg.n_heads} kv={cfg.n_kv_heads}")
print(f"  K={K} block_size={BLOCK_SIZE} top_k={TOP_K} n_blocks={N_BLOCKS}")
print(f"  params: {params['grand_total']/1e6:.1f}M ({params_gb:.2f} GB)")
print(f"  KV cache: {kv_gb:.2f} GB ({CTX:,} tokens)")
print(f"  K_bar: {kbar_gb:.4f} GB ({N_BLOCKS} blocks)")
print(f"  Estimated total VRAM: {est_total:.2f} GB")

if est_total > vram_free * 0.9:
    print(f"  WARNING: est {est_total:.1f}GB > 90% of free {vram_free:.1f}GB — aborting")
    sys.exit(1)

# --- Build model ---
print("\nBuilding model...")
m = build_model(cfg, dev)
m.precompute_medusa_tokens()
print(f"  model VRAM: {torch.cuda.memory_allocated()/1e9:.2f} GB")

# --- Allocate block-structured KV cache ---
print("Allocating block-sparse KV cache...")
cache_k = [torch.zeros(1, N_BLOCKS, BLOCK_SIZE, cfg.n_kv_heads, cfg.head_dim,
                        device=dev, dtype=dt) for _ in range(cfg.n_layers)]
cache_v = [torch.zeros(1, N_BLOCKS, BLOCK_SIZE, cfg.n_kv_heads, cfg.head_dim,
                        device=dev, dtype=dt) for _ in range(cfg.n_layers)]
k_bar = [torch.zeros(1, N_BLOCKS, cfg.n_kv_heads, cfg.head_dim,
                     device=dev, dtype=dt) for _ in range(cfg.n_layers)]
print(f"  cache VRAM: {torch.cuda.memory_allocated()/1e9:.2f} GB")

# --- Fill cache with random data (simulating 20M tokens already processed) ---
print(f"\nFilling cache to {CTX:,} tokens...")
n_filled = CTX // BLOCK_SIZE  # 4883 complete blocks
current_fill = 0  # current block is empty (next step writes here)

# Fill all complete blocks with random K,V
for lk, lv, lb in zip(cache_k, cache_v, k_bar):
    lk[:, :n_filled].normal_(0, 0.1)
    lv[:, :n_filled].normal_(0, 0.1)
    # Compute k_bar for each filled block
    lb[:, :n_filled] = lk[:, :n_filled].mean(dim=2)  # mean over block_size dim

print(f"  VRAM after fill: {torch.cuda.memory_allocated()/1e9:.2f} GB")
print(f"  filled blocks: {n_filled}, current_fill: {current_fill}")

# --- Warmup ---
input_ids = torch.randint(0, cfg.vocab_size, (1, K1), device=dev)
print("\nWarmup...")
for _ in range(3):
    with torch.no_grad():
        h = m.forward_moba(input_ids, cache_k, cache_v, k_bar,
                           n_filled, current_fill)
        tokens = m.medusa_argmax_fast(h[:, -1])
    torch.cuda.synchronize()

# --- Measure ---
print("Measuring...")
times = []
N_STEPS = 10
for step in range(N_STEPS):
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    with torch.no_grad():
        h = m.forward_moba(input_ids, cache_k, cache_v, k_bar,
                           n_filled, current_fill)
        tokens = m.medusa_argmax_fast(h[:, -1])
    torch.cuda.synchronize()
    times.append(time.perf_counter() - t0)

times.sort()
med = times[len(times) // 2]
ms = med * 1000
tps = K1 / med
us = ms / K1 * 1000

print(f"\n{'='*60}")
print(f"MoBA at {CTX:,} context (block_size={BLOCK_SIZE}, top_k={TOP_K})")
print(f"  {K1} tokens/step")
print(f"  step: {ms:.1f} ms")
print(f"  speed: {tps:,.0f} tok/s")
print(f"  per-token: {us:.2f} us/tok")
print(f"  peak VRAM: {torch.cuda.max_memory_allocated()/1e9:.2f} GB")
print(f"  target: 50,000 tok/s  {'PASS' if tps >= 50000 else 'FAIL'}")
print(f"{'='*60}")

# --- Sweep block_size and top_k ---
print(f"\n{'='*60}")
print("Sweep: block_size x top_k")
print(f"{'BS':>6} {'k':>3} {'step(ms)':>9} {'tok/s':>10} {'us/tok':>8}")
print("-" * 40)

for bs in [4097, 8192, 16384, 32768]:
    for tk in [1, 2, 3, 5, 8]:
        n_blk = CTX // bs
        # Reallocate cache for this block size
        ck = [torch.zeros(1, n_blk, bs, cfg.n_kv_heads, cfg.head_dim,
                          device=dev, dtype=dt) for _ in range(cfg.n_layers)]
        cv = [torch.zeros(1, n_blk, bs, cfg.n_kv_heads, cfg.head_dim,
                          device=dev, dtype=dt) for _ in range(cfg.n_layers)]
        kb = [torch.zeros(1, n_blk, cfg.n_kv_heads, cfg.head_dim,
                          device=dev, dtype=dt) for _ in range(cfg.n_layers)]

        nf = n_blk - 1  # leave room for current block
        for lk, lv, lb in zip(ck, cv, kb):
            lk[:, :nf].normal_(0, 0.1)
            lv[:, :nf].normal_(0, 0.1)
            lb[:, :nf] = lk[:, :nf].mean(dim=2)

        # Update model's MoBA params
        m.blocks[0].attn.block_size = bs
        m.blocks[0].attn.top_k = tk

        # Warmup
        for _ in range(2):
            with torch.no_grad():
                h = m.forward_moba(input_ids, ck, cv, kb, nf, 0)
                _ = m.medusa_argmax_fast(h[:, -1])
        torch.cuda.synchronize()

        # Measure
        ts = []
        for _ in range(5):
            torch.cuda.synchronize()
            t0 = time.perf_counter()
            with torch.no_grad():
                h = m.forward_moba(input_ids, ck, cv, kb, nf, 0)
                _ = m.medusa_argmax_fast(h[:, -1])
            torch.cuda.synchronize()
            ts.append(time.perf_counter() - t0)
        ts.sort()
        med = ts[len(ts) // 2]
        ms = med * 1000
        tps = K1 / med
        us = ms / K1 * 1000
        print(f"{bs:6d} {tk:3d} {ms:9.1f} {tps:10,.0f} {us:8.2f}")

        del ck, cv, kb
        torch.cuda.empty_cache()

# Restore default
m.blocks[0].attn.block_size = BLOCK_SIZE
m.blocks[0].attn.top_k = TOP_K