#!/usr/bin/env python3
# =============================================================================
#  Project   : localvm-research
#  File      : experiments/candidate_01/benchmark_scale.py
#  Purpose   : Scale run — 32B model whose q8 does NOT fit beside the resident
#              base: q4 resident + layer-streamed q8 verification sweeps
#  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 scale benchmark (the regime the architecture exists for).

Qwen3-32B on a 48 GB Mac: q4 (17.5 GB) resident; q8 (34.8 GB) cannot be
co-resident — sweeps stream it layer-by-layer from SSD (StreamingVerifier).
Baseline: pure q4 (the only real alternative on this machine). Quality judged
by Qwen3-8B-bf16 (independent judge; the 32B bf16 obviously cannot run).

Usage:
    .venv/bin/python benchmark_scale.py [--per-domain 2] [--max-tokens 96]
"""

from __future__ import annotations

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

import mlx.core as mx
from huggingface_hub import snapshot_download
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
from streaming_verifier import StreamingVerifier  # noqa: E402

Q4_REPO = "mlx-community/Qwen3-32B-4bit"
Q8_REPO = "mlx-community/Qwen3-32B-8bit"
JUDGE_REPO = "mlx-community/Qwen3-8B-bf16"


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):
        nxt = int(mx.argmax(model(inp, cache=cache)[0, -1]).item())
        if nxt == tokenizer.eos_token_id:
            break
        tokens.append(nxt)
        inp = mx.array([[nxt]])
    return tokens, time.perf_counter() - t0


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--per-domain", type=int, default=2)
    ap.add_argument("--max-tokens", type=int, default=96)
    ap.add_argument("--window", type=int, default=32)
    ap.add_argument("--taus", default="2.0")
    ap.add_argument("--modes", default="margin,verify-all")
    args = ap.parse_args()

    q4_path = snapshot_download(Q4_REPO)
    q8_path = snapshot_download(Q8_REPO)
    domains = json.loads((REPO_ROOT / "benchmarks/datasets/eval_prompts.json").read_text())["domains"]

    print("loading q4 resident …", flush=True)
    base_model, tokenizer = load(q4_path)
    verifier = StreamingVerifier(q8_path)
    q8_bytes = verifier.weight_bytes
    print(f"q8 checkpoint (streamed): {q8_bytes/1e9:.1f} GB", flush=True)

    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)})

    print("baseline: pure q4 …", flush=True)
    q4_out, q4_times = [], []
    for k, p in enumerate(prompts):
        toks, dt = greedy_baseline(base_model, tokenizer, p["ids"], args.max_tokens)
        q4_out.append(toks); q4_times.append((len(toks), dt))
        print(f"  {k+1}/{len(prompts)} ({len(toks)} tok, {len(toks)/dt:.1f} tok/s)", flush=True)

    outputs = {"pure_q4": q4_out}
    runs = []
    for mode in args.modes.split(","):
        for tau in ([float(x) for x in args.taus.split(",")] if mode == "margin" else [2.0]):
            print(f"runtime: mode={mode} tau={tau} W={args.window} …", flush=True)
            outs, agg = [], {"tokens": 0, "deferred": 0, "sweeps": 0, "rollbacks": 0,
                             "sweep_s": 0.0, "gen_s": 0.0, "logical_bytes": 0, "io_s": []}
            for k, p in enumerate(prompts):
                toks, st = generate_deferred(
                    base_model, verifier, tokenizer, p["ids"],
                    args.max_tokens, tau, args.window, mode, q8_bytes)
                outs.append(toks)
                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
                print(f"  {k+1}/{len(prompts)} ({st.tokens_out} tok, {st.sweeps} sweeps, "
                      f"{st.rollbacks} rollbacks, last sweep io {verifier.last_sweep_io_s:.1f}s)",
                      flush=True)
            n = max(agg["tokens"], 1)
            runs.append({
                "mode": mode, "tau": tau, "window": args.window,
                "tokens_per_s": n / (agg["gen_s"] + agg["sweep_s"]),
                "deferral_rate": agg["deferred"] / n,
                "rollback_rate": agg["rollbacks"] / n,
                "sweep_latency_s_mean": agg["sweep_s"] / max(agg["sweeps"], 1),
                "logical_verify_bytes_per_token": agg["logical_bytes"] / n,
                "raw": {k: v for k, v in agg.items() if k != "io_s"},
            })
            outputs[f"{mode}_tau{tau}"] = outs
            r = runs[-1]
            print(f"  tok/s={r['tokens_per_s']:.2f} sweepLat={r['sweep_latency_s_mean']:.1f}s "
                  f"GB/token(logical)={r['logical_verify_bytes_per_token']/1e9:.2f}", flush=True)

    print("freeing 32B models; loading 8B bf16 judge …", flush=True)
    del base_model, verifier
    gc.collect(); mx.clear_cache()
    judge, _ = load(JUDGE_REPO)
    quality = {}
    for name, outs in outputs.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)
            lp = sel - mx.logsumexp(sel, axis=-1, keepdims=True)
            tok_lp = mx.take_along_axis(lp, mx.array(toks)[:, None], axis=-1)
            mx.eval(tok_lp)
            vals.append(float(mx.mean(tok_lp).item()))
        quality[name] = {"mean_logprob_8b_judge": sum(vals) / len(vals), "n": len(vals)}
        print(f"  {name:>18}: {quality[name]['mean_logprob_8b_judge']:.4f}", flush=True)

    ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    out_dir = REPO_ROOT / "results" / "candidate_01_scale32b" / ts
    out_dir.mkdir(parents=True)
    (out_dir / "results.json").write_text(json.dumps({
        "experiment": "candidate_01_scale32b",
        "author": "Simon-Pierre Boucher",
        "contact": "contact@spboucher.ai",
        "manifest": collect_manifest(),
        "config": vars(args),
        "models": {"base": Q4_REPO, "verify": Q8_REPO, "judge": JUDGE_REPO},
        "q8_streamed_bytes": q8_bytes,
        "baseline_pure_q4_tokens_per_s":
            sum(t for t, _ in q4_times) / max(sum(d for _, d in q4_times), 1e-9),
        "runs": runs,
        "quality_8b_judge": quality,
    }, indent=2))
    print(f"\nwrote {out_dir / 'results.json'}")


if __name__ == "__main__":
    main()