#!/usr/bin/env python3
"""Negative results: settings that looked promising and did not pay off.

Two knobs are measured end to end under the same 8 GiB ceiling as the headline:

prefetch
    Paging the next layer's experts while the current layer computes. The pager
    can predict them from the previous token's routing, but only a small share
    of the predicted experts are actually used, so the extra reads are mostly
    waste. ``SDQ_PREFETCH_LAYER_AHEAD=1`` turns it on.

threads
    Only n_expert_used = 8 experts per token can run in parallel inside one
    layer, so more threads mostly contend for the same rows.

Output: bench/results/strict_prefetch.json and bench/results/strict_threads.json
"""

from __future__ import annotations

import argparse
import json
import os
import re
import statistics
import subprocess
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
REPO = HERE.parent

PROMPTS = [
    "Explain in one sentence why SSDs are slow.",
    "請用一句話介紹你自己。",
    "Write a short paragraph about SSD caching.",
]


def run(strict8g: Path, binary: str, model: str, prompt: str, threads: int,
        n_predict: int, ctx: int, ubatch: int, budget_mb: int,
        record: Path, reserve_mb: int, env_extra: dict) -> dict:
    cmd = [
        sys.executable, str(strict8g),
        "--limit-mb", str(budget_mb), "--record", str(record), "--",
        binary, "-m", model, "-p", prompt, "-n", str(n_predict),
        "-c", str(ctx), "-t", str(threads), "-b", str(ubatch),
        "--temp", "0", "--no-chat", "--no-interactive",
        "--ram-budget-mb", str(budget_mb),
        "--stats-file", str(record.with_suffix(".stats")),
    ]
    env = dict(os.environ, SDQ_RESERVE_MB=str(reserve_mb))
    env.update(env_extra)
    p = subprocess.run(cmd, capture_output=True, text=True, env=env)
    out = p.stdout + p.stderr

    rec, stats = {}, {}
    if record.exists():
        try:
            rec = json.loads(record.read_text())
        except Exception:
            pass
    sp = record.with_suffix(".stats")
    if sp.exists():
        for line in sp.read_text().splitlines():
            if line.strip().startswith("{"):
                stats = json.loads(line)

    m = re.search(r"= ([\d.]+) tok/s", out)
    return {
        "rc": p.returncode,
        "strict_ok": bool(rec.get("strict_ok", False)),
        "tps": float(m.group(1)) if m else None,
        "hit_pct": stats.get("hit_rate"),
        "mb_per_token": (stats["ssd_bytes_per_token"] / 1048576.0) if stats.get("ssd_bytes_per_token") else None,
        "prefetch_issued": stats.get("prefetch_issued"),
        "prefetch_used": stats.get("prefetch_used"),
        "peak_anon_mb": rec.get("peak_anon_mb"),
        "peak_swap_mb": rec.get("peak_swap_mb"),
    }


def sweep(arms: dict, a, label: str, out_name: str) -> None:
    strict8g = REPO / "tools" / "strict8g.py"
    (HERE / "results").mkdir(parents=True, exist_ok=True)
    rows = []
    for r in range(a.rounds):
        for name, cfg in arms.items():
            prompt = PROMPTS[r % len(PROMPTS)]
            rec = HERE / "results" / f".{out_name}-{name}-{r}.json"
            row = run(strict8g, a.bin, a.model, prompt, cfg["threads"],
                      a.n_predict, a.ctx, a.ubatch, a.budget_mb, rec,
                      cfg["reserve_mb"], cfg.get("env", {}))
            row["arm"] = name
            row["round"] = r
            rows.append(row)
            print(f"[{r + 1}/{a.rounds}] {name:<14} tps={row['tps']}  "
                  f"hit={row['hit_pct']}  MB/tok={row['mb_per_token']}  "
                  f"ok={row['strict_ok']}", flush=True)
            rec.unlink(missing_ok=True)
            Path(str(rec) + ".stats").with_suffix(".stats").unlink(missing_ok=True)

    summary = {}
    for name in arms:
        tps = [r["tps"] for r in rows if r["arm"] == name and r["tps"]]
        summary[name] = {
            "n": len(tps),
            "median_tps": statistics.median(tps) if tps else None,
            "mean_tps": statistics.mean(tps) if tps else None,
            "stdev_tps": statistics.stdev(tps) if len(tps) > 1 else None,
            "min_tps": min(tps) if tps else None,
            "max_tps": max(tps) if tps else None,
            "all_strict_ok": all(r["strict_ok"] for r in rows if r["arm"] == name),
        }
    out = HERE / "results" / out_name
    out.write_text(json.dumps({"label": label, "rows": rows, "summary": summary},
                              indent=2, ensure_ascii=False))
    print(json.dumps(summary, indent=2, ensure_ascii=False))
    print(f"-> {out}")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--bin", required=True)
    ap.add_argument("--which", choices=["prefetch", "threads"], required=True)
    ap.add_argument("--rounds", type=int, default=6)
    ap.add_argument("--budget-mb", type=int, default=8192)
    ap.add_argument("--reserve-mb", type=int, default=512)
    ap.add_argument("--ctx", type=int, default=4096)
    ap.add_argument("--ubatch", type=int, default=128)
    ap.add_argument("--n-predict", type=int, default=48)
    a = ap.parse_args()

    if a.which == "prefetch":
        arms = {
            "prefetch_off": {"threads": 8,
                             "reserve_mb": a.reserve_mb, "env": {}},
            "prefetch_layer_ahead": {"threads": 8, "reserve_mb": a.reserve_mb,
                                     "env": {"SDQ_PREFETCH_LAYER_AHEAD": "1"}},
            "prefetch_end_token": {"threads": 8, "reserve_mb": a.reserve_mb,
                                   "env": {"SDQ_PREFETCH": "1",
                                           "SDQ_PREFETCH_LAYER_AHEAD": "0"}},
        }
        sweep(arms, a, "prefetch", "strict_prefetch.json")
    else:
        arms = {
            f"threads_{t}": {"threads": t, "reserve_mb": a.reserve_mb, "env": {}}
            for t in (4, 8, 10, 12, 16)
        }
        sweep(arms, a, "threads", "strict_threads.json")


if __name__ == "__main__":
    main()