"""Find the 500M base config that keeps ~1M tok/s at 20M context.

Physics:
  - Weight reads: (base + medusa) * 2 bytes / bandwidth
  - Attention: 4 * top_k * L * (K+1) * block_size * d_model FLOPs
  - tok/s = (K+1) / step_time

  500M base needs large d and/or many L, but attention scales with d*L.
  Sweet spot: K where weight_read_time ≈ attention_time.
"""
import time
import torch
from config import Config
from model import build_model

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

CTX = 20_000_000
TOP_K = 3

# Candidate 500M base configs
# Key constraint: KV cache at 20M must fit: 20M * n_kv * hd * 4 * L < 40 GB
#   => n_kv * hd * L < 500
configs = [
    # (name, d, L, H, kv, ffn_mult, K, block_size, top_k)
    # Current baseline (for comparison)
    ("baseline_228M",  384, 1, 12, 4, 2.667, 4096, 4097, 3),

    # 500M base configs — varying d, L, K
    # d=1024, L=8, kv=2: KV=20.5GB, base~481M
    ("500M_d1024_L8_K4096", 1024, 8, 16, 2, 16.0, 4096, 4097, 3),
    ("500M_d1024_L8_K2048", 1024, 8, 16, 2, 16.0, 2048, 2049, 3),
    ("500M_d1024_L8_K1024", 1024, 8, 16, 2, 16.0, 1024, 1025, 3),
    ("500M_d1024_L8_K512",  1024, 8, 16, 2, 16.0, 512,  513,  3),

    # d=1024, L=8, kv=2, top_k=1 (faster attention, less context)
    ("500M_d1024_L8_K4096_t1", 1024, 8, 16, 2, 16.0, 4096, 4097, 1),
    ("500M_d1024_L8_K2048_t1", 1024, 8, 16, 2, 16.0, 2048, 2049, 1),
    ("500M_d1024_L8_K1024_t1", 1024, 8, 16, 2, 16.0, 1024, 1025, 1),

    # d=768, L=12, kv=2: KV=18.4GB, base~277M (smaller base, more layers)
    ("277M_d768_L12_K4096", 768, 12, 12, 2, 10.0, 4096, 4097, 3),
    ("277M_d768_L12_K2048", 768, 12, 12, 2, 10.0, 2048, 2049, 3),
    ("277M_d768_L12_K1024", 768, 12, 12, 2, 10.0, 1024, 1025, 3),

    # d=1536, L=4, kv=2: KV=15.7GB, base~498M (larger d, fewer layers)
    ("500M_d1536_L4_K4096", 1536, 4, 24, 2, 4.0, 4096, 4097, 3),
    ("500M_d1536_L4_K2048", 1536, 4, 24, 2, 4.0, 2048, 2049, 3),
    ("500M_d1536_L4_K1024", 1536, 4, 24, 2, 4.0, 1024, 1025, 3),
    ("500M_d1536_L4_K4096_t1", 1536, 4, 24, 2, 4.0, 4096, 4097, 1),
]

print(f"\n{'name':>28} {'d':>5} {'L':>3} {'K':>5} {'tk':>3} "
      f"{'base(M)':>8} {'total(M)':>9} {'KV(GB)':>7} "
      f"{'step(ms)':>9} {'tok/s':>10} {'us/tok':>8}")
print("-" * 95)

for name, d, L, H, kv, ffn, K, BS, tk in configs:
    cfg = Config(
        d_model=d, n_layers=L, n_heads=H, n_kv_heads=kv,
        ffn_mult=ffn, medusa_heads=K, medusa_rank=1, medusa_hidden=0,
        window_size=0, max_seq_len=CTX + K + 100, use_alibi=False,
        moba_block_size=BS, moba_top_k=tk, dtype="bfloat16",
    )

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

    params = cfg.count_params()
    base_m = params["base_total"] / 1e6
    total_m = params["grand_total"] / 1e6
    kv_gb = CTX * kv * (d // H) * 4 * L / 1e9

    if kv_gb + total_m * 2 / 1e9 > vram_free * 0.9:
        print(f"{name:>28} {d:5d} {L:3d} {K:5d} {tk:3d} "
              f"{base_m:8.1f} {total_m:9.1f} {kv_gb:7.2f}  --- VRAM too high ---")
        continue

    n_blocks = (CTX // BS) + 2
    try:
        m = build_model(cfg, dev)
        m.precompute_medusa_tokens()

        cache_k = [torch.zeros(1, n_blocks, BS, kv, d // H,
                                device=dev, dtype=dt) for _ in range(L)]
        cache_v = [torch.zeros(1, n_blocks, BS, kv, d // H,
                                device=dev, dtype=dt) for _ in range(L)]
        k_bar = [torch.zeros(1, n_blocks, kv, d // H,
                              device=dev, dtype=dt) for _ in range(L)]

        # Fill cache
        nf = CTX // BS
        for lk, lv, lb in zip(cache_k, cache_v, k_bar):
            lk[:, :nf].normal_(0, 0.1)
            lv[:, :nf].normal_(0, 0.1)
            lb[:, :nf] = lk[:, :nf].mean(dim=2)

        input_ids = torch.randint(0, cfg.vocab_size, (1, K + 1), device=dev)

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

        # Measure
        times = []
        for _ in range(5):
            torch.cuda.synchronize()
            t0 = time.perf_counter()
            with torch.no_grad():
                h = m.forward_moba(input_ids, cache_k, cache_v, k_bar, nf, 0)
                _ = 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 = (K + 1) / med
        us = ms / (K + 1) * 1000

        print(f"{name:>28} {d:5d} {L:3d} {K:5d} {tk:3d} "
              f"{base_m:8.1f} {total_m:9.1f} {kv_gb:7.2f} "
              f"{ms:9.1f} {tps:10,.0f} {us:8.2f}")

    except Exception as e:
        print(f"{name:>28} {d:5d} {L:3d} {K:5d} {tk:3d} "
              f"{base_m:8.1f} {total_m:9.1f} {kv_gb:7.2f}  --- ERROR: {e} ---")

    torch.cuda.empty_cache()
    try:
        del m, cache_k, cache_v, k_bar
    except:
        pass
    torch.cuda.empty_cache()

print(f"\n  peak VRAM: {torch.cuda.max_memory_allocated()/1e9:.2f} GB")