"""INT4-vs-NVFP4 KV-cache quantization sweep on Laguna-XS.2.

For each layer's real K and V activations, measures reconstruction RMSE across
the grid  {nvfp4, int4} × {headdim-block, per-channel-block} × {absmax, mse}.

The baseline cell (nvfp4 / headdim / absmax) is what vLLM's NVFP4 KV kernel ships
today. Every cell has identical memory (4-bit data + one scale per 16 elements),
so this isolates reconstruction quality. K and V are reported separately to
expose any KIVI-style asymmetry (per-channel K, per-token V).

Usage:
    python -m scripts.quant_sweep [--max-new 200] [--n-alphas 32]
"""
from __future__ import annotations

import argparse
import sys

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache

sys.path.insert(0, "/home/alex/poolside-hackathon-kv-quant")
from kv_quant import BLOCK, SWEEP_CELLS, SWEEP_BASELINE, sweep_tensor

MODEL = "poolside/Laguna-XS.2"
PROMPT = (
    "Solve step by step. A train leaves city A at 60 km/h. Two hours later a second "
    "train leaves the same station on the same track at 90 km/h. How many hours after "
    "the second train departs will it catch up to the first train? Show your reasoning."
)


def _cell_label(cell) -> str:
    return f"{cell[0]:<6} {cell[1]:<8} {cell[2]:<6}"


def _print_table(title: str, abs_rmse: dict, ratio: dict) -> None:
    base = SWEEP_BASELINE
    print(f"\n{title}")
    print(f"  {'format':<6} {'layout':<8} {'calib':<6} {'avg RMSE':>10} {'vs vLLM':>10}")
    print(f"  {'-'*6} {'-'*8} {'-'*6} {'-'*10} {'-'*10}")
    for c in [base] + [c for c in SWEEP_CELLS if c != base]:
        tag = "baseline" if c == base else f"{(1.0 - ratio[c]) * 100.0:+.1f}%"
        print(f"  {_cell_label(c)} {abs_rmse[c]:>10.5f} {tag:>10}")
    best = min(SWEEP_CELLS, key=lambda c: ratio[c])
    print(f"  best: {best[0]}/{best[1]}/{best[2]}  "
          f"({(1.0 - ratio[best]) * 100.0:+.1f}% vs baseline)")


def run_sweep(per_layer_kv, n_alphas, device):
    """per_layer_kv: list of (K, V) [n_kv, seq, head_dim]. Returns means + ratios."""
    base = SWEEP_BASELINE
    k_abs = {c: [] for c in SWEEP_CELLS}
    k_rat = {c: [] for c in SWEEP_CELLS}
    v_abs = {c: [] for c in SWEEP_CELLS}
    v_rat = {c: [] for c in SWEEP_CELLS}
    n_used = 0
    for K, V in per_layer_kv:
        if K.shape[1] < BLOCK:
            continue
        n_used += 1
        ks = sweep_tensor(K.to(device), n_alphas)
        vs = sweep_tensor(V.to(device), n_alphas)
        for c in SWEEP_CELLS:
            k_abs[c].append(ks[c]); k_rat[c].append(ks[c] / max(ks[base], 1e-12))
            v_abs[c].append(vs[c]); v_rat[c].append(vs[c] / max(vs[base], 1e-12))
    mean = lambda d: {c: sum(x) / max(len(x), 1) for c, x in d.items()}
    return n_used, mean(k_abs), mean(k_rat), mean(v_abs), mean(v_rat)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--max-new", type=int, default=200)
    ap.add_argument("--n-alphas", type=int, default=32)
    args = ap.parse_args()

    print(f"[load] {MODEL} ...", flush=True)
    tok = AutoTokenizer.from_pretrained(MODEL)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL, torch_dtype=torch.bfloat16, device_map="auto")
    model.eval()
    device = next(model.parameters()).device
    cfg = model.config
    print(f"[load] layers={cfg.num_hidden_layers} kv_heads={cfg.num_key_value_heads} "
          f"head_dim={getattr(cfg, 'head_dim', None)} device={device}", flush=True)

    msgs = [{"role": "user", "content": PROMPT}]
    input_ids = tok.apply_chat_template(
        msgs, add_generation_prompt=True, return_tensors="pt", return_dict=False,
    ).to(device)

    cache = DynamicCache()
    with torch.no_grad():
        model.generate(input_ids, max_new_tokens=args.max_new,
                        past_key_values=cache, use_cache=True, do_sample=False)

    per_layer = [(layer.keys[0].float().cpu(), layer.values[0].float().cpu())
                 for layer in cache.layers]
    seq = per_layer[0][0].shape[1] if per_layer else 0
    print(f"[seq] {seq} tokens cached across {len(per_layer)} layers", flush=True)

    n_used, k_abs, k_rat, v_abs, v_rat = run_sweep(per_layer, args.n_alphas, device)

    print(f"\n{'='*56}")
    print(f"INT4 vs NVFP4 KV sweep  —  {n_used} layers, seq cropped to /{BLOCK}")
    print("All cells: identical memory (4-bit + 1 scale per 16 elems, ~3.56x vs BF16).")
    _print_table("KEY cache (avg RMSE over layers):", k_abs, k_rat)
    _print_table("VALUE cache (avg RMSE over layers):", v_abs, v_rat)


if __name__ == "__main__":
    main()