#!/usr/bin/env python3
"""Analyse cache-size sweep results.

For each subdir under --run-root containing engine_step.jsonl:
  - read all per-step records
  - bin by cache_size
  - report median/p90 of step_duration_us, build_meta_us per bin
  - fit step_duration_us ~ a + b * cache_size  (linear least squares)
  - tabulate connector tax(cache_size) = mc_step - plain_step (if plain present)
  - optionally render matplotlib plot if matplotlib available

Outputs:
  results/<date>/SUMMARY.md       human report
  results/<date>/per_config.json  machine-readable
  results/<date>/figure.png       (optional)
"""

import argparse
import json
import statistics
from pathlib import Path


def load_steps(p: Path):
    rows = []
    with open(p) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                rows.append(json.loads(line))
            except Exception:
                pass
    return rows


def percentile(xs, p):
    if not xs:
        return None
    xs = sorted(xs)
    k = max(0, min(len(xs) - 1, int(p / 100.0 * (len(xs) - 1))))
    return xs[k]


def linfit(xs, ys):
    """Tiny linear least squares. Returns (slope, intercept)."""
    n = len(xs)
    if n < 2:
        return None, None
    sx = sum(xs); sy = sum(ys); sxx = sum(x*x for x in xs)
    sxy = sum(x*y for x, y in zip(xs, ys))
    denom = n * sxx - sx * sx
    if denom == 0:
        return None, None
    b = (n * sxy - sx * sy) / denom
    a = (sy - b * sx) / n
    return b, a


def bucket(rows, key="cache_size", n_bins=10):
    """Equal-width bin on cache_size; returns dict bin_id -> list of rows."""
    if not rows:
        return {}
    vmax = max(r.get(key, 0) for r in rows)
    if vmax <= 0:
        return {}
    width = vmax / n_bins
    out: dict[int, list[dict]] = {}
    for r in rows:
        v = r.get(key, 0)
        bid = min(n_bins - 1, max(0, int(v / width))) if width > 0 else 0
        out.setdefault(bid, []).append(r)
    return out, width


def analyse_config(cfg_name: str, cfg_dir: Path):
    eng_path = cfg_dir / "engine_step.jsonl"
    if not eng_path.exists() or eng_path.stat().st_size == 0:
        return None
    raw = load_steps(eng_path)
    if not raw:
        return None

    # Filter: skip first 500 steps (cold start), and steps with no cache_size.
    base = [r for r in raw[500:] if r.get("cache_size", -1) >= 0]
    if not base:
        return None

    # Decode-only filter: steps where the scheduler did NOT touch any
    # new/resumed request and total tokens == n_running_total (each running
    # request emits exactly one token). This gives the cleanest per-step
    # baseline since prefill chunks dominate step time at high token counts.
    decode_only = [r for r in base
                   if r.get("prefill_tokens", 0) == 0
                   and r.get("decode_tokens", 0) > 0]
    # Fall back to "all post-warmup" if decode-only is too sparse
    rows = decode_only if len(decode_only) >= 200 else base
    decode_share = len(decode_only) / max(1, len(base))

    cache_max = max(r["cache_size"] for r in rows)
    bins, width = bucket(rows, n_bins=10)

    per_bin = []
    for bid in sorted(bins):
        rs = bins[bid]
        sd = [r["step_duration_us"] for r in rs if "step_duration_us" in r]
        bm = [r.get("build_meta_us", 0) for r in rs]
        cs = [r["cache_size"] for r in rs]
        per_bin.append({
            "bin_id": bid,
            "cache_size_mid": (bid + 0.5) * width,
            "n": len(rs),
            "cache_size_p50": percentile(cs, 50),
            "step_duration_us_p50": percentile(sd, 50),
            "step_duration_us_p90": percentile(sd, 90),
            "build_meta_us_p50": percentile(bm, 50),
            "build_meta_us_p90": percentile(bm, 90),
        })

    # Fit per-step duration vs cache size on all records (not bin averages)
    sd_b, sd_a = linfit([r["cache_size"] for r in rows if "step_duration_us" in r],
                        [r["step_duration_us"] for r in rows if "step_duration_us" in r])
    bm_b, bm_a = linfit([r["cache_size"] for r in rows if "build_meta_us" in r],
                        [r.get("build_meta_us", 0) for r in rows if "build_meta_us" in r])

    # Worker-side timings if available. Filename is `worker_step.r<rank>.jsonl`
    # because os.path.splitext keeps the .jsonl extension.
    worker_path = None
    for c in sorted(cfg_dir.glob("worker_step.r*.jsonl")):
        worker_path = c
        break
    if worker_path is None:
        worker_path = cfg_dir / "missing"
    worker_summary = None
    if worker_path.exists() and worker_path.stat().st_size > 0:
        wrows = load_steps(worker_path)
        if wrows:
            gf = [r["get_finished_us"] for r in wrows if "get_finished_us" in r]
            sl = [r["start_load_kv_us"] for r in wrows if "start_load_kv_us" in r]
            worker_summary = {
                "n": len(wrows),
                "get_finished_us_p50": percentile(gf, 50),
                "get_finished_us_p90": percentile(gf, 90),
                "get_finished_us_p99": percentile(gf, 99),
                "start_load_kv_us_p50": percentile(sl, 50),
                "start_load_kv_us_p90": percentile(sl, 90),
            }

    return {
        "config": cfg_name,
        "n_steps_total": len(raw),
        "n_steps_after_warmup": len(base),
        "n_steps_decode_only": len(decode_only),
        "decode_share": decode_share,
        "rows_used_for_fit": "decode_only" if rows is decode_only else "all_post_warmup",
        "cache_size_max": cache_max,
        "per_bin": per_bin,
        "fit_step_duration": {"slope_us_per_block": sd_b, "intercept_us": sd_a},
        "fit_build_meta":    {"slope_us_per_block": bm_b, "intercept_us": bm_a},
        "worker_summary": worker_summary,
    }


def render(root: Path, all_cfg: dict):
    lines = ["# Cache-size sweep — summary\n"]
    lines.append(f"Run root: `{root}`\n")
    lines.append("## Per-config fit (`step_duration_us ≈ a + b · cache_size`)\n")
    lines.append("| config | n steps | cache max | step_dur p50 (μs) | build_meta p50 (μs) | slope (μs / 1k blocks) | intercept (μs) |")
    lines.append("|---|---:|---:|---:|---:|---:|---:|")
    for cfg, r in all_cfg.items():
        if r is None:
            lines.append(f"| {cfg} | — | — | — | — | — | — |")
            continue
        last_bin = r["per_bin"][-1] if r["per_bin"] else {}
        slope = r["fit_step_duration"]["slope_us_per_block"]
        intercept = r["fit_step_duration"]["intercept_us"]
        slope1k = (slope or 0) * 1000
        lines.append(
            f"| {cfg} | {r['n_steps_after_warmup']} | {r['cache_size_max']} | "
            f"{last_bin.get('step_duration_us_p50','-')} | "
            f"{last_bin.get('build_meta_us_p50','-')} | "
            f"{slope1k:.1f} | {intercept:.1f} |"
            if slope is not None else
            f"| {cfg} | {r['n_steps_after_warmup']} | {r['cache_size_max']} | - | - | - | - |"
        )

    # Per-bin tables
    for cfg, r in all_cfg.items():
        if r is None:
            continue
        lines.append(f"\n### {cfg} — per-bin\n")
        lines.append("| bin | cache mid | n | step_dur p50 | step_dur p90 | build_meta p50 | build_meta p90 |")
        lines.append("|---:|---:|---:|---:|---:|---:|---:|")
        for b in r["per_bin"]:
            lines.append(
                f"| {b['bin_id']} | {b['cache_size_mid']:.0f} | {b['n']} | "
                f"{b['step_duration_us_p50']} | {b['step_duration_us_p90']} | "
                f"{b['build_meta_us_p50']} | {b['build_meta_us_p90']} |"
            )
        if r["worker_summary"]:
            w = r["worker_summary"]
            lines.append(f"\n*worker side (n={w['n']})* — "
                         f"get_finished p50/p90/p99 = "
                         f"{w['get_finished_us_p50']}/{w['get_finished_us_p90']}/"
                         f"{w['get_finished_us_p99']} μs; "
                         f"start_load_kv p50/p90 = "
                         f"{w['start_load_kv_us_p50']}/{w['start_load_kv_us_p90']} μs\n")

    # Tax vs cache for mc vs plain
    plain = all_cfg.get("plain")
    mc = all_cfg.get("mooncake_both")
    mc_dr = all_cfg.get("mooncake_both_drfix")
    noop = all_cfg.get("noop_connector")

    def _tax_table(label, baseline, target):
        if not (baseline and target):
            return
        lines.append(f"\n## {label}\n")
        lines.append("| bin | cache mid | baseline step p50 | target step p50 | tax (μs) | tax (%) |")
        lines.append("|---:|---:|---:|---:|---:|---:|")
        for bp, bm in zip(baseline["per_bin"], target["per_bin"]):
            if bp["step_duration_us_p50"] and bm["step_duration_us_p50"]:
                tax = bm["step_duration_us_p50"] - bp["step_duration_us_p50"]
                pct = tax / bp["step_duration_us_p50"] * 100
                lines.append(
                    f"| {bp['bin_id']} | {bp['cache_size_mid']:.0f} | "
                    f"{bp['step_duration_us_p50']} | {bm['step_duration_us_p50']} | "
                    f"{tax:+d} | {pct:+.1f} |"
                )

    _tax_table("Connector tax(cache_size) — mooncake_both vs plain", plain, mc)
    _tax_table("Connector tax(cache_size) — mooncake_both_drfix vs plain", plain, mc_dr)
    _tax_table("DR-fix savings — mooncake_both vs mooncake_both_drfix",
               mc_dr, mc)  # baseline=fixed, target=original → "tax" = savings

    if plain and noop:
        # Framework cost: noop_connector tax = pure dispatch
        lines.append("\n## Framework cost — noop_connector vs plain\n")
        lines.append("| bin | cache mid | plain step p50 | noop step p50 | tax (μs) |")
        lines.append("|---:|---:|---:|---:|---:|")
        for bp, bn in zip(plain["per_bin"], noop["per_bin"]):
            if bp["step_duration_us_p50"] and bn["step_duration_us_p50"]:
                tax = bn["step_duration_us_p50"] - bp["step_duration_us_p50"]
                lines.append(
                    f"| {bp['bin_id']} | {bp['cache_size_mid']:.0f} | "
                    f"{bp['step_duration_us_p50']} | {bn['step_duration_us_p50']} | "
                    f"{tax:+d} |"
                )

    out_md = root / "SUMMARY.md"
    out_md.write_text("\n".join(lines) + "\n")
    out_json = root / "per_config.json"
    out_json.write_text(json.dumps(all_cfg, indent=2, default=str))
    print(f"  wrote {out_md}")
    print(f"  wrote {out_json}")

    # Optional plot
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5))
        colors = {"plain": "tab:blue", "noop_connector": "tab:orange",
                  "mooncake_both": "tab:red",
                  "mooncake_both_drfix": "tab:green"}
        for cfg, r in all_cfg.items():
            if r is None: continue
            xs = [b["cache_size_mid"] for b in r["per_bin"]]
            ys = [b["step_duration_us_p50"] or 0 for b in r["per_bin"]]
            zs = [b["build_meta_us_p50"] or 0 for b in r["per_bin"]]
            c = colors.get(cfg, None)
            ax1.plot(xs, ys, marker="o", label=cfg, color=c)
            ax2.plot(xs, zs, marker="s", label=cfg, color=c)
        ax1.set_xlabel("cache_size (blocks)")
        ax1.set_ylabel("step_duration_us p50")
        ax1.set_title("Per-step scheduler time vs |cache|")
        ax1.legend(); ax1.grid(True, alpha=0.3)
        ax2.set_xlabel("cache_size (blocks)")
        ax2.set_ylabel("build_meta_us p50")
        ax2.set_title("build_connector_meta time vs |cache|")
        ax2.legend(); ax2.grid(True, alpha=0.3)
        fig.tight_layout()
        fig.savefig(root / "figure.png", dpi=120)
        print(f"  wrote {root/'figure.png'}")
    except Exception as e:
        print(f"  (skipped plot: {e})")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--run-root", type=Path, required=True)
    args = ap.parse_args()

    cfgs = {}
    for d in sorted(args.run_root.iterdir()):
        if not d.is_dir(): continue
        r = analyse_config(d.name, d)
        cfgs[d.name] = r
        if r:
            sl = r["fit_step_duration"]["slope_us_per_block"]
            print(f"  {d.name}: n={r['n_steps_after_warmup']} "
                  f"cache_max={r['cache_size_max']} "
                  f"slope={(sl or 0)*1000:.2f} μs/1k blocks")
        else:
            print(f"  {d.name}: no data")
    render(args.run_root, cfgs)


if __name__ == "__main__":
    main()