#!/usr/bin/env python3
# =============================================================================
#  Project   : localvm-research
#  File      : experiments/candidate_01/benchmark.py
#  Purpose   : End-to-end evaluation of the margin-gated deferred-refinement
#              runtime vs pure-q4 / pure-q8 baselines
#  Author    : Simon-Pierre Boucher
#  Contact   : contact@spboucher.ai
#  Created   : 2026-08-12
#  Modified  : 2026-08-12
#  Platform  : macOS / Apple Silicon (arm64) — MLX / Metal
#  License   : All rights reserved (research code)
# =============================================================================
"""Candidate-01 benchmark.

Build once (downloads + quantizes):
    .venv/bin/python benchmark.py --build
Run:
    .venv/bin/python benchmark.py [--per-domain 4] [--max-tokens 128]
        [--window 32] [--taus 1.0,2.0]
"""

from __future__ import annotations

import argparse
import difflib
import json
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

import mlx.core as mx
from mlx_lm import load

REPO_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO_ROOT / "benchmarks"))
sys.path.insert(0, str(Path(__file__).parent / "implementation"))
from hardware_manifest import collect_manifest  # noqa: E402
from runtime import generate_deferred  # noqa: E402

MODELS_DIR = Path(__file__).parent / "implementation" / "models"
HF_MODEL = "mlx-community/Qwen3-1.7B-bf16"


def build() -> None:
    from mlx_lm import convert

    for bits in (4, 8):
        out = MODELS_DIR / f"q{bits}"
        if out.exists():
            print(f"{out} exists, skipping")
            continue
        print(f"converting {HF_MODEL} → q{bits} …", flush=True)
        convert(HF_MODEL, mlx_path=str(out), quantize=True, q_bits=bits, q_group_size=64)
    print("build done")


def dir_weight_bytes(d: Path) -> int:
    return sum(f.stat().st_size for f in d.glob("*.safetensors"))


def greedy_baseline(model, tokenizer, prompt_ids, max_tokens):
    from mlx_lm.models.cache import make_prompt_cache

    cache = make_prompt_cache(model)
    tokens = []
    inp = mx.array(list(prompt_ids))[None]
    t0 = time.perf_counter()
    for _ in range(max_tokens):
        logits = model(inp, cache=cache)
        nxt = int(mx.argmax(logits[0, -1]).item())
        if nxt == tokenizer.eos_token_id:
            break
        tokens.append(nxt)
        inp = mx.array([[nxt]])
    return tokens, time.perf_counter() - t0


def judge_outputs(outputs_by_config: dict, prompts: list[dict]) -> dict:
    """Quality-level metric: mean per-token logprob of each config's generated
    continuation under the bf16 reference model (higher = better). Token-exact
    fidelity is incoherent on Metal (1.56%/token prefill/decode flips), so the
    judge scores usefulness of the text the system actually produced."""
    import gc

    gc.collect(); mx.clear_cache()
    judge, _ = load(HF_MODEL)
    scores = {}
    for name, outs in outputs_by_config.items():
        vals = []
        for p, toks in zip(prompts, outs):
            if len(toks) < 2:
                continue
            full = p["ids"] + list(toks)
            logits = judge(mx.array(full)[None])[0]
            sel = logits[len(p["ids"]) - 1 : len(full) - 1].astype(mx.float32)
            logprobs = sel - mx.logsumexp(sel, axis=-1, keepdims=True)
            idx = mx.array(toks)
            tok_lp = mx.take_along_axis(logprobs, idx[:, None], axis=-1)
            mx.eval(tok_lp)
            vals.append(float(mx.mean(tok_lp).item()))
        scores[name] = {"mean_logprob_bf16": sum(vals) / len(vals), "n": len(vals)}
    del judge
    gc.collect(); mx.clear_cache()
    return scores


def fidelity(a: list[int], b: list[int]) -> float:
    """Similarity of two token sequences (difflib ratio — robust to length
    drift after divergence)."""
    if not a and not b:
        return 1.0
    return difflib.SequenceMatcher(None, a, b).ratio()


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--build", action="store_true")
    ap.add_argument("--per-domain", type=int, default=4)
    ap.add_argument("--max-tokens", type=int, default=128)
    ap.add_argument("--window", type=int, default=32)
    ap.add_argument("--taus", default="1.0,2.0")
    args = ap.parse_args()
    if args.build:
        build()
        return

    domains = json.loads((REPO_ROOT / "benchmarks/datasets/eval_prompts.json").read_text())["domains"]
    q4_dir, q8_dir = MODELS_DIR / "q4", MODELS_DIR / "q8"
    q8_bytes = dir_weight_bytes(q8_dir)
    q4_bytes = dir_weight_bytes(q4_dir)
    print(f"resident q4: {q4_bytes/1e9:.2f} GB · streamed q8: {q8_bytes/1e9:.2f} GB", flush=True)

    base_model, tokenizer = load(str(q4_dir))
    verify_model, _ = load(str(q8_dir))

    prompts = []
    for domain, plist in domains.items():
        for prompt in plist[: args.per_domain]:
            ids = tokenizer.apply_chat_template(
                [{"role": "user", "content": prompt}], add_generation_prompt=True)
            prompts.append({"domain": domain, "ids": list(ids)})

    # baselines
    print("baseline: pure q8 greedy …", flush=True)
    q8_out, q8_times = [], []
    for p in prompts:
        toks, dt = greedy_baseline(verify_model, tokenizer, p["ids"], args.max_tokens)
        q8_out.append(toks); q8_times.append((len(toks), dt))
    print("baseline: pure q4 greedy …", flush=True)
    q4_out, q4_times = [], []
    for p in prompts:
        toks, dt = greedy_baseline(base_model, tokenizer, p["ids"], args.max_tokens)
        q4_out.append(toks); q4_times.append((len(toks), dt))

    def toks_per_s(times):
        n = sum(t for t, _ in times); s = sum(d for _, d in times)
        return n / s if s else 0.0

    configs = []
    for mode in ("margin", "verify-all"):
        for tau in ([float(x) for x in args.taus.split(",")] if mode == "margin" else [2.0]):
            configs.append({"mode": mode, "tau": tau})

    outputs_by_config = {"pure_q4": q4_out, "pure_q8": q8_out}
    results = []
    for cfg in configs:
        print(f"runtime: mode={cfg['mode']} tau={cfg['tau']} W={args.window} …", flush=True)
        fid, agg = [], {"tokens": 0, "deferred": 0, "sweeps": 0, "rollbacks": 0,
                        "sweep_s": 0.0, "gen_s": 0.0, "logical_bytes": 0}
        cfg_outputs = []
        for p, ref in zip(prompts, q8_out):
            toks, st = generate_deferred(
                base_model, verify_model, tokenizer, p["ids"],
                args.max_tokens, cfg["tau"], args.window, cfg["mode"], q8_bytes)
            cfg_outputs.append(toks)
            fid.append(fidelity(toks, ref))
            agg["tokens"] += st.tokens_out; agg["deferred"] += st.deferred
            agg["sweeps"] += st.sweeps; agg["rollbacks"] += st.rollbacks
            agg["sweep_s"] += st.sweep_time_s; agg["gen_s"] += st.gen_time_s
            agg["logical_bytes"] += st.sweep_logical_bytes
        n = max(agg["tokens"], 1)
        results.append({
            **cfg, "window": args.window,
            "fidelity_vs_q8_mean": sum(fid) / len(fid),
            "tokens_per_s": n / (agg["gen_s"] + agg["sweep_s"]),
            "deferral_rate": agg["deferred"] / n,
            "rollback_rate": agg["rollbacks"] / n,
            "sweeps_per_100tok": 100 * agg["sweeps"] / n,
            "sweep_latency_s_mean": agg["sweep_s"] / max(agg["sweeps"], 1),
            "logical_verify_bytes_per_token": agg["logical_bytes"] / n,
            "raw": agg,
        })
        outputs_by_config[f"{cfg['mode']}_tau{cfg['tau']}"] = cfg_outputs
        r = results[-1]
        print(f"  fidelity={r['fidelity_vs_q8_mean']:.4f} tok/s={r['tokens_per_s']:.1f} "
              f"defer={r['deferral_rate']:.2f} rollback={r['rollback_rate']:.3f} "
              f"MB/token(logical)={r['logical_verify_bytes_per_token']/1e6:.0f}", flush=True)

    print("judging outputs with bf16 reference …", flush=True)
    quality = judge_outputs(outputs_by_config, prompts)
    for name, s in quality.items():
        print(f"  {name:>18}: mean logprob (bf16 judge) = {s['mean_logprob_bf16']:.4f}", flush=True)

    payload = {
        "experiment": "candidate_01_deferred_refinement",
        "author": "Simon-Pierre Boucher",
        "contact": "contact@spboucher.ai",
        "manifest": collect_manifest(),
        "config": vars(args),
        "model": HF_MODEL,
        "q4_resident_bytes": q4_bytes, "q8_stream_bytes": q8_bytes,
        "baselines": {
            "pure_q4": {"tokens_per_s": toks_per_s(q4_times),
                        "fidelity_vs_q8_mean": sum(fidelity(a, b) for a, b in zip(q4_out, q8_out)) / len(q8_out)},
            "pure_q8": {"tokens_per_s": toks_per_s(q8_times), "fidelity_vs_q8_mean": 1.0},
        },
        "runs": results,
        "quality_bf16_judge": quality,
    }
    ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    out_dir = REPO_ROOT / "results" / "candidate_01" / ts
    out_dir.mkdir(parents=True)
    (out_dir / "results.json").write_text(json.dumps(payload, indent=2))
    print(f"\nwrote {out_dir / 'results.json'}")
    print("baselines:", json.dumps(payload["baselines"], indent=1))


if __name__ == "__main__":
    main()