"""
turbo_benchmark.py — faithful TurboQuant baseline (Zandieh et al., ICLR 2026).

Haar rotation + per-coordinate Lloyd-Max codebooks + optional QJL residual for
unbiased inner-product estimation. Separate from benchmark.py; writes
artifacts/turbo_codebooks.pt.

Usage:
    python benchmark.py --stage dump   # or use existing calib_caches.pt
    python turbo_benchmark.py --stage fit
    python turbo_benchmark.py --stage cheap
"""

from __future__ import annotations

import argparse
import json
import os
import time

import torch

from vqkv.metrics import (key_cosine, cache_mse, inner_product_distortion,
                          attention_output, attn_output_cosine, attn_output_error)

ARTIFACT_DIR = os.environ.get("VQKV_ARTIFACTS", "./artifacts")
MODEL_ID = os.environ.get("LAGUNA_ID", "poolside/Laguna-XS.2")
CALIB_DATASET = os.environ.get("CALIB_DATASET", "SWE-bench/SWE-smith-trajectories")
CALIB_SPLIT = os.environ.get("CALIB_SPLIT", "tool")
EVAL_SOURCE = os.environ.get("EVAL_SOURCE", "swesmith")

_HOTPOTQA_PROMPT = (
    "Answer the question based on the given passages. "
    "Only give me the answer and do not output any other words.\n\n"
    "The following are given passages.\n{context}\n\n"
    "Answer the question based on the given passages. "
    "Only give me the answer and do not output any other words.\n\n"
    "Question: {input}\nAnswer:"
)


# ============================================================================
# Lloyd-Max 1-D scalar quantizer (the MSE-optimal levels for a given sample)
# ============================================================================
def lloyd_max_1d(samples: torch.Tensor, n_levels: int, iters: int = 30,
                 seed: int = 0) -> torch.Tensor:
    """Optimal scalar (Lloyd-Max) codebook for a 1-D distribution given samples.

    samples: (M,) 1-D values drawn from the (concentrated) coordinate dist.
    returns: (n_levels,) sorted reconstruction levels.

    This is 1-D k-means; for the rotated unit-norm coordinates the distribution
    is the same across coordinates in high dim, so one shared codebook per
    bit-width suffices (matching the paper's single precomputed Beta codebook).
    """
    s = samples.flatten()
    s = s[torch.isfinite(s)]
    # init levels at quantiles so empty bins are rare
    qs = torch.linspace(0.0, 1.0, n_levels + 2, device=s.device)[1:-1]
    levels = torch.quantile(s, qs)
    for _ in range(iters):
        # assign each sample to nearest level
        idx = torch.bucketize(s, (levels[1:] + levels[:-1]) / 2)
        new = levels.clone()
        for j in range(n_levels):
            m = idx == j
            if m.any():
                new[j] = s[m].mean()
        shift = (new - levels).abs().max()
        levels = new
        if shift < 1e-6:
            break
    return torch.sort(levels).values


def _haar_rotation(d: int, seed: int = 0, device=None) -> torch.Tensor:
    g = torch.Generator().manual_seed(seed)
    a = torch.randn(d, d, generator=g)
    q, r = torch.linalg.qr(a)
    # sign-correct so Q is Haar-distributed (QR sign ambiguity)
    q = q * torch.sign(torch.diag(r)).unsqueeze(0)
    if device is not None:
        q = q.to(device)
    return q


# ============================================================================
# TurboQuant-MSE : rotation -> per-coordinate Lloyd-Max -> norm rescale
# ============================================================================
class TurboQuantMSE:
    def __init__(self, nbits: int = 3, seed: int = 0):
        self.nbits = nbits
        self.seed = seed
        self._rot = None       # (d, d)
        self._levels = None    # (K,) shared Lloyd-Max levels

    def fit(self, calib: torch.Tensor):
        """calib: (N, d). Fit rotation + shared per-coordinate level set."""
        d = calib.shape[-1]
        self._rot = _haar_rotation(d, self.seed, calib.device)
        xn = calib / calib.norm(dim=-1, keepdim=True).clamp_min(1e-8)
        r = xn @ self._rot                       # (N, d) rotated unit-norm coords
        # all coordinates share the same concentrated dist -> pool them
        pool = r.flatten()
        if pool.numel() > 2_000_000:             # cap for speed
            pool = pool[torch.randperm(pool.numel(), device=pool.device)[:2_000_000]]
        self._levels = lloyd_max_1d(pool, 1 << self.nbits)
        return self

    def _quantize(self, x):
        norms = x.norm(dim=-1, keepdim=True).clamp_min(1e-8)
        xn = x / norms
        y = xn @ self._rot
        edges = (self._levels[1:] + self._levels[:-1]) / 2
        idx = torch.bucketize(y, edges).clamp(0, self._levels.numel() - 1)
        return idx, norms

    def roundtrip(self, x):
        idx, norms = self._quantize(x)
        y_hat = self._levels[idx]
        x_hat = y_hat @ self._rot.T
        return x_hat * norms

    def to(self, device):
        self._rot = self._rot.to(device)
        self._levels = self._levels.to(device)
        return self


# ============================================================================
# QJL inner-product estimator (faithful, ASYMMETRIC) -- NOT a reconstruction.
#
# Important: QJL does not reconstruct a vector. It estimates <q, k> directly,
# with the query JL-transformed but UNQUANTIZED and only the key residual
# sign-quantized (the asymmetric estimator of Zandieh et al.). It therefore
# cannot be expressed as roundtrip(x)->x_hat without reintroducing the very
# bias it removes. We expose it as a standalone estimator used only in the
# inner-product-bias evaluation, not in the cosine/MSE reconstruction metrics.
#
#   <q, k> ~= <q, k_hat_mse>  +  (sqrt(pi/2) * ||r|| / m) * sum_i (g_i . q) * sign(g_i . r)
#
# where r = k_n - k_hat_mse_n is the residual in unit-norm space, g_i are the
# rows of a gaussian JL matrix, and ||r|| is the stored residual norm.
# ============================================================================
class QJLResidualIP:
    def __init__(self, mse: "TurboQuantMSE", qjl_rows: int = None, seed: int = 0):
        self.mse = mse
        self.seed = seed
        self.qjl_rows = qjl_rows
        self._G = None

    def fit(self, calib: torch.Tensor):
        d = calib.shape[-1]
        m = self.qjl_rows or d
        g = torch.Generator().manual_seed(self.seed + 7)
        self._G = torch.randn(m, d, generator=g).to(calib.device)
        return self

    def estimate_ip(self, q, k):
        """Unbiased estimate of rowwise <q, k> using MSE stage + QJL residual.
        q, k: (N, d). Returns (N,) inner-product estimates."""
        m = self._G.shape[0]
        # stage-1 reconstruction in unit-norm space
        knorm = k.norm(dim=-1, keepdim=True).clamp_min(1e-8)
        kn = k / knorm
        idx, _ = self.mse._quantize(k)
        kn_hat = (self.mse._levels[idx]) @ self.mse._rot.T
        score1 = (q * (kn_hat * knorm)).sum(-1)           # <q, k_hat_mse>
        # stage-2 residual correction (asymmetric: q unquantized, residual signed)
        r = (kn - kn_hat)
        rnorm = r.norm(dim=-1, keepdim=True).clamp_min(1e-12)
        rn = r / rnorm
        qg = q @ self._G.T                                # (N, m)  query JL, unquantized
        rs = torch.sign(rn @ self._G.T)                   # (N, m)  residual sign bits
        score2 = ((torch.pi / 2) ** 0.5 / m) * (qg * rs).sum(-1) * rnorm.squeeze(-1) * knorm.squeeze(-1)
        return score1 + score2

    def to(self, device):
        self.mse.to(device)
        self._G = self._G.to(device)
        return self


# ============================================================================
# Config table. Each TurboQuant variant uses TurboQuant-MSE for RECONSTRUCTION
# (the metric that feeds cosine/MSE). For keys we ALSO build a QJL inner-product
# estimator (the faithful Prod second stage), reported only in the ip_bias
# column since QJL is an IP estimator, not a reconstruction.
#
# `qjl=True` marks variants that add the QJL key estimator. 1-bit is included
# as the aggressive floor; at 1 bit the MSE stage is effectively sign, so the
# QJL residual carries most of the IP fidelity there.
# ============================================================================
def turbo_configs(bits_list):
    cfgs = []
    for b in bits_list:
        cfgs.append((f"turbo-mse-{b}b", b, False))
        cfgs.append((f"turbo-prod-{b}b (K:+qjl)", b, True))
    return cfgs


# ============================================================================
# STAGE: fit -- per-layer codebooks from calib_caches.pt
# ============================================================================
def stage_fit(bits_list):
    path = os.path.join(ARTIFACT_DIR, "calib_caches.pt")
    if not os.path.exists(path):
        raise FileNotFoundError(
            f"{path} not found. Run `python benchmark.py --stage dump` first "
            f"to produce the calibration caches this script reuses.")
    blob = torch.load(path, weights_only=False)
    calib, meta = blob["calib"], blob["meta"]
    hd = meta["head_dim"]
    dev = "cuda" if torch.cuda.is_available() else "cpu"

    fitted = {}
    for name, b, use_qjl in turbo_configs(bits_list):
        per_layer = {}
        t0 = time.time()
        for i, c in calib.items():
            kf = c["k"].reshape(-1, hd)[:200_000].to(dev)
            vf = c["v"].reshape(-1, hd)[:200_000].to(dev)
            kq = TurboQuantMSE(nbits=b).fit(kf)
            vq = TurboQuantMSE(nbits=b).fit(vf)
            qjl = QJLResidualIP(kq).fit(kf) if use_qjl else None
            per_layer[i] = {"kq": kq, "vq": vq, "qjl": qjl}
        fitted[name] = per_layer
        print(f"[fit] {name}: {len(per_layer)} layer-codebooks in {time.time()-t0:.1f}s")
    torch.save({"fitted": fitted, "meta": meta, "bits": bits_list},
               os.path.join(ARTIFACT_DIR, "turbo_codebooks.pt"))
    print(f"[fit] saved -> {ARTIFACT_DIR}/turbo_codebooks.pt")


# ============================================================================
# Trace flattening -- minimal copy (kept independent of benchmark.py)
# ============================================================================
def flatten_trace(example, tok) -> str:
    raw = (example.get("messages") or example.get("trajectory")
           or example.get("conversations"))
    if raw is None:
        return json.dumps(example)[:200_000]
    if isinstance(raw, str):
        try:
            raw = json.loads(raw)
        except json.JSONDecodeError:
            return raw[:200_000]
    norm = []
    for m in raw:
        role = m.get("role") or m.get("from") or "user"
        content = m.get("content") or m.get("value") or ""
        if isinstance(content, list):
            content = "\n".join(
                it.get("text", str(it)) if isinstance(it, dict) else str(it)
                for it in content)
        tc = m.get("tool_calls")
        if tc:
            tct = json.dumps(tc, ensure_ascii=False)
            content = (content + "\n" + tct).strip() if content else tct
        role = {"human": "user", "gpt": "assistant", "tool": "user"}.get(role, role)
        if not content.strip():
            continue
        norm.append({"role": role, "content": content})
    merged = []
    for m in norm:
        if merged and merged[-1]["role"] == m["role"]:
            merged[-1]["content"] += "\n\n" + m["content"]
        else:
            merged.append(dict(m))
    try:
        return tok.apply_chat_template(merged, tokenize=False,
                                       add_generation_prompt=False)
    except Exception:
        return "\n\n".join(f"{m['role']}: {m['content']}" for m in merged)


def flatten_longbench(example) -> str:
    """Format a LongBench hotpotqa example (context + input) as a plain string."""
    return _HOTPOTQA_PROMPT.format(
        context=example["context"], input=example["input"])


def _load_longbench_hotpotqa():
    """Load THUDM/LongBench hotpotqa, bypassing the deprecated dataset script."""
    from datasets import load_dataset as _ld
    for fname in ("hotpotqa_e.jsonl", "hotpotqa.jsonl"):
        try:
            return _ld(
                "json",
                data_files=f"hf://datasets/THUDM/LongBench/data/{fname}",
                split="train",
            )
        except Exception:
            continue
    return _ld("THUDM/LongBench", name="hotpotqa", split="test")


def _load_cheap_eval_dataset(n_eval: int, eval_source: str, tok):
    """Return (dataset_slice, get_text, max_len, label) for stage_cheap."""
    from datasets import load_dataset

    if eval_source == "longbench-hotpotqa":
        ds = _load_longbench_hotpotqa()
        # First n_eval rows — disjoint from dump calib (last n_calib rows).
        end = min(n_eval, len(ds))
        ds = ds.select(range(0, end))
        label = f"LongBench hotpotqa (rows 0–{end - 1})"
        return ds, flatten_longbench, 32768, label

    ds = load_dataset(CALIB_DATASET, split=CALIB_SPLIT)
    start = 500  # held-out offset (matches benchmark.py --stage cheap)
    end = min(start + n_eval, len(ds))
    ds = ds.select(range(start, end))
    label = f"{CALIB_DATASET} split={CALIB_SPLIT} (rows {start}–{end - 1})"
    return ds, lambda ex: flatten_trace(ex, tok), 16384, label


def load_model_and_meta():
    from transformers import AutoModelForCausalLM, AutoTokenizer
    tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID, dtype=torch.bfloat16, device_map="cuda", trust_remote_code=True)
    model.eval()
    cfg = model.config
    full = [i for i, t in enumerate(cfg.layer_types) if t == "full_attention"]
    meta = {"full_layers": full, "n_kv_heads": cfg.num_key_value_heads,
            "n_q_heads": cfg.num_attention_heads,
            "head_dim": cfg.head_dim, "n_layers": cfg.num_hidden_layers}
    print(f"[meta] full-attention layers ({len(full)}): {full}")
    return model, tok, meta


# ============================================================================
# STAGE: cheap -- full metric set on held-out traces (mirrors benchmark.py)
# ============================================================================
def stage_cheap(n_eval=64, max_len=None, eval_source: str | None = None,
                min_len=2048):
    import collections
    from tqdm import tqdm
    from transformers.cache_utils import DynamicCache

    eval_source = eval_source or EVAL_SOURCE
    blob = torch.load(os.path.join(ARTIFACT_DIR, "turbo_codebooks.pt"),
                      weights_only=False)
    fitted, meta = blob["fitted"], blob["meta"]
    full = meta["full_layers"]
    hd   = meta["head_dim"]
    n_q  = meta.get("n_q_heads", 48)   # fallback for turbo_codebooks.pt without this field

    # Window for O(T²) attention metrics (matches benchmark.py)
    ATTN_WIN = 512

    # Lookup: config name -> (nbits, use_qjl) for bits_per_elt reporting
    cfg_lookup = {name: (b, use_qjl) for name, b, use_qjl in turbo_configs(blob["bits"])}

    model, tok, _ = load_model_and_meta()
    dev = model.device
    if str(dev) == "meta":
        dev = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    ds, get_text, default_max_len, source_label = _load_cheap_eval_dataset(
        n_eval, eval_source, tok)
    if max_len is None:
        max_len = default_max_len
    print(f"[turbo-cheap] eval_source={eval_source}  {source_label}  "
          f"rows={len(ds)}  max_len={max_len}")

    for per_layer in fitted.values():
        for entry in per_layer.values():
            entry["kq"].to(dev); entry["vq"].to(dev)
            if entry["qjl"] is not None:
                entry["qjl"].to(dev)

    class EvalDump(DynamicCache):
        def __init__(self):
            super().__init__(); self.d = {i: {} for i in full}
        def update(self, ks, vs, li, ck=None):
            if li in set(full):
                self.d[li]["k"] = ks.detach()[0].permute(1, 0, 2).float()
                self.d[li]["v"] = vs.detach()[0].permute(1, 0, 2).float()
            return super().update(ks, vs, li, ck)

    trace_rows = []
    n_used = 0

    for ex in tqdm(ds, desc=f"turbo-cheap/{eval_source}"):
        text = get_text(ex)
        ids = tok(text, return_tensors="pt", truncation=True,
                  max_length=max_len).to(dev)
        if ids["input_ids"].shape[1] < min_len:
            continue
        cache = EvalDump()
        with torch.no_grad():
            model.model(**ids, past_key_values=cache, use_cache=True)
        n_used += 1

        # Synthetic Q: generated once per (trace, layer), reused across configs.
        # Using synthetic Q for the QJL estimator is more faithful than the
        # previous keys-as-proxy approach: q0[qi] are proper query vectors,
        # kr0[ki] are the keys being estimated.
        synth_q = {}
        for i in full:
            s = cache.d[i]["k"].shape[0]
            win = min(s, ATTN_WIN)
            q_rand = torch.randn(win, n_q, hd, device=dev)
            synth_q[i] = q_rand / q_rand.norm(dim=-1, keepdim=True).clamp_min(1e-8)

        for name, per_layer in fitted.items():
            acc = collections.defaultdict(float)
            nL = 0
            for i in full:
                k = cache.d[i]["k"]; v = cache.d[i]["v"]   # (s, h, d)
                s, h, d = k.shape
                e = per_layer[i]
                k_hat = e["kq"].roundtrip(k.reshape(-1, d)).reshape(s, h, d)
                v_hat = e["vq"].roundtrip(v.reshape(-1, d)).reshape(s, h, d)

                acc["key_cos"] += key_cosine(k, k_hat)
                acc["val_cos"] += key_cosine(v, v_hat)
                acc["key_mse"] += cache_mse(k, k_hat)
                acc["val_mse"] += cache_mse(v, v_hat)

                # Windowed attention/IP metrics on last ATTN_WIN tokens
                win = min(s, ATTN_WIN)
                kw, kw_hat = k[-win:], k_hat[-win:]
                vw, vw_hat = v[-win:], v_hat[-win:]
                q_syn = synth_q[i]                          # (win, n_q, d)

                out_ref, _ = attention_output(q_syn, kw, vw, n_q)
                out_hat, _ = attention_output(q_syn, kw_hat, vw_hat, n_q)
                acc["attn_cos"]          += attn_output_cosine(out_ref, out_hat)
                acc["attn_output_error"] += attn_output_error(out_ref, out_hat)

                # IP metrics. For Prod configs, use the faithful QJL asymmetric
                # estimator (designed for unbiased <q,k> estimation). For MSE
                # configs, use the plain reconstructed inner product.
                q0  = q_syn[:, 0, :]    # (win, d) — head 0 of synthetic Q
                kr0 = kw[:, 0, :]       # (win, d) — head 0 of reference K
                kh0 = kw_hat[:, 0, :]   # (win, d) — head 0 of reconstructed K
                qi  = torch.randint(0, win, (4096,), device=dev)
                ki  = torch.randint(0, win, (4096,), device=dev)
                ip_ref = (q0[qi] * kr0[ki]).sum(-1)
                if e["qjl"] is not None:
                    ip_hat = e["qjl"].estimate_ip(q0[qi], kr0[ki])
                else:
                    ip_hat = (q0[qi] * kh0[ki]).sum(-1)
                acc["ip_bias"] += (ip_hat - ip_ref).mean().item()
                acc["ip_rel"]  += ((ip_hat - ip_ref).abs() /
                                   ip_ref.abs().clamp_min(1e-6)).mean().item()
                nL += 1

            trace_rows.append({
                "trace_len":         ids["input_ids"].shape[1],
                "config":            name,
                "key_cos":           acc["key_cos"]           / nL,
                "val_cos":           acc["val_cos"]           / nL,
                "key_mse":           acc["key_mse"]           / nL,
                "val_mse":           acc["val_mse"]           / nL,
                "attn_cos":          acc["attn_cos"]          / nL,
                "attn_output_error": acc["attn_output_error"] / nL,
                "ip_rel":            acc["ip_rel"]            / nL,
                "ip_bias":           acc["ip_bias"]           / nL,
            })

    # Aggregate across traces: one summary row per config
    agg = collections.defaultdict(lambda: collections.defaultdict(list))
    for r in trace_rows:
        for col in ("key_cos", "val_cos", "key_mse", "val_mse",
                    "attn_cos", "attn_output_error", "ip_rel", "ip_bias"):
            agg[r["config"]][col].append(r[col])

    COLS = ("key_cos", "val_cos", "key_mse", "val_mse",
            "attn_cos", "attn_output_error", "ip_rel", "ip_bias")

    summary = []
    for name in cfg_lookup:
        if name not in agg:
            continue
        b, use_qjl = cfg_lookup[name]
        # bpe: nbits/coord + 1 bit/coord for QJL signs (m=d rows) + fp16 norm
        bpe = b + (1 if use_qjl else 0) + 16.0 / hd
        cols = agg[name]
        n = len(cols["key_cos"])
        row = {"config": name, "bits_per_elt": round(bpe, 4),
               "n_traces": n, "eval_source": eval_source}
        for col in COLS:
            row[col] = round(sum(cols[col]) / n, 5)
        summary.append(row)

    print(f"\n[turbo-cheap] mean metrics over {n_used} traces ({eval_source}):")
    print(f"  {'config':30s}  {'bpe':>5}  {'key_cos':>8}  {'val_cos':>8}  "
          f"{'key_mse':>9}  {'val_mse':>9}  {'attn_cos':>9}  {'attn_err':>9}  "
          f"{'ip_rel':>8}  {'ip_bias':>9}")
    for row in summary:
        print(f"  {row['config']:30s}  {row['bits_per_elt']:5.2f}  "
              f"{row['key_cos']:8.4f}  {row['val_cos']:8.4f}  "
              f"{row['key_mse']:9.5f}  {row['val_mse']:9.5f}  "
              f"{row['attn_cos']:9.4f}  {row['attn_output_error']:9.4f}  "
              f"{row['ip_rel']:8.5f}  {row['ip_bias']:9.6f}")

    out_name = ("turbo_cheap_metrics_hotpotqa.json"
                if eval_source == "longbench-hotpotqa"
                else "turbo_cheap_metrics.json")
    out_path = os.path.join(ARTIFACT_DIR, out_name)
    json.dump(summary, open(out_path, "w"), indent=2)
    print(f"[turbo-cheap] saved -> {out_path}")


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--stage", required=True, choices=["fit", "cheap"])
    ap.add_argument("--bits", type=int, nargs="+", default=[4, 2, 1],
                    help="bit-widths to sweep (default: 4 2 1)")
    ap.add_argument("--n_eval", type=int, default=64)
    ap.add_argument(
        "--eval_source", type=str, default=None,
        choices=["swesmith", "longbench-hotpotqa"],
        help="cheap stage: eval corpus (default: EVAL_SOURCE env or swesmith)")
    ap.add_argument("--max_len", type=int, default=None,
                    help="cheap stage: max prompt tokens (default: 16384 swesmith, "
                         "32768 longbench-hotpotqa)")
    args = ap.parse_args()
    if args.stage == "fit":
        stage_fit(args.bits)
    elif args.stage == "cheap":
        stage_cheap(n_eval=args.n_eval, max_len=args.max_len,
                    eval_source=args.eval_source)


if __name__ == "__main__":
    main()