#!/usr/bin/env python3
"""B3 5-policy re-test analyser.

Compute TTFT/TPOT/E2E mean/p50/p90/p99 for each policy from
metrics.jsonl, compare against the historical b3_policy_comparison.json
that drives fig_b3_latency_bars.png, and emit a side-by-side table
plus a new figure with the same layout as the original.

Usage:
  python analyze_b3_replay.py --root <outroot> [--old-data <path>] [--figure <path>]
"""

import argparse
import json
import statistics
from pathlib import Path


POLICIES = ["lmetric", "load_only", "sticky", "unified", "unified_v2"]


def pct(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 summarise(path):
    rows = [json.loads(l) for l in open(path) if l.strip()]
    ok = [r for r in rows if not r.get("error")]
    ttft = [r["ttft_s"] * 1000 for r in ok if r.get("ttft_s") is not None]
    tpot = [r["tpot_s"] * 1000 for r in ok if r.get("tpot_s")]
    e2e = [r["latency_s"] * 1000 for r in ok if r.get("latency_s") is not None]
    return {
        "n_total": len(rows),
        "n_ok": len(ok),
        "ttft_mean_ms": statistics.mean(ttft) if ttft else None,
        "ttft_p50_ms": pct(ttft, 50),
        "ttft_p90_ms": pct(ttft, 90),
        "ttft_p99_ms": pct(ttft, 99),
        "tpot_mean_ms": statistics.mean(tpot) if tpot else None,
        "tpot_p50_ms": pct(tpot, 50),
        "tpot_p90_ms": pct(tpot, 90),
        "tpot_p99_ms": pct(tpot, 99),
        "e2e_mean_ms": statistics.mean(e2e) if e2e else None,
        "e2e_p50_ms": pct(e2e, 50),
        "e2e_p90_ms": pct(e2e, 90),
        "e2e_p99_ms": pct(e2e, 99),
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", type=Path, required=True)
    ap.add_argument("--old-data", type=Path,
                    default=Path("analysis/characterization/window_1_results/b3_policy_comparison.json"))
    ap.add_argument("--figure", type=Path, default=None)
    args = ap.parse_args()

    new = {}
    for p in POLICIES:
        path = args.root / p / "metrics.jsonl"
        if not path.exists():
            print(f"MISSING: {path}")
            continue
        new[p] = summarise(path)

    old = {}
    if args.old_data.exists():
        d = json.load(open(args.old_data))
        for r in d.get("rows", []):
            old[r["policy"]] = {
                "ttft_p50_ms": r["ttft_p50_s"] * 1000,
                "ttft_p90_ms": r["ttft_p90_s"] * 1000,
                "ttft_p99_ms": r["ttft_p99_s"] * 1000,
                "tpot_p90_ms": r["tpot_p90_s"] * 1000,
                "e2e_p90_ms":  r.get("e2e_p90_s", 0) * 1000,
            }

    def fmt(v): return f"{v:.0f}" if v is not None else "-"
    def pctd(a, b):
        if a is None or b is None or a == 0: return "-"
        return f"{(b/a-1)*100:+.1f}%"

    # Headline table
    print(f"\n# NEW: today's re-test")
    print(f"{'policy':<14}{'n_ok':>6}{'TTFTp50':>10}{'TTFTp90':>10}{'TTFTp99':>10}{'TPOTp90':>10}{'E2Ep90':>10}")
    print("-" * 70)
    for p in POLICIES:
        if p not in new: continue
        r = new[p]
        print(f"{p:<14}{r['n_ok']:>6}{fmt(r['ttft_p50_ms']):>9}ms{fmt(r['ttft_p90_ms']):>9}ms{fmt(r['ttft_p99_ms']):>9}ms{fmt(r['tpot_p90_ms']):>9}ms{fmt(r['e2e_p90_ms']):>9}ms")

    print(f"\n# OLD: window_1_results/b3_policy_comparison.json")
    print(f"{'policy':<14}{'TTFTp50':>10}{'TTFTp90':>10}{'TTFTp99':>10}{'TPOTp90':>10}{'E2Ep90':>10}")
    print("-" * 60)
    for p in POLICIES:
        if p not in old: continue
        r = old[p]
        print(f"{p:<14}{fmt(r['ttft_p50_ms']):>9}ms{fmt(r['ttft_p90_ms']):>9}ms{fmt(r['ttft_p99_ms']):>9}ms{fmt(r['tpot_p90_ms']):>9}ms{fmt(r['e2e_p90_ms']):>9}ms")

    print(f"\n# DRIFT: today vs old (same policy)")
    print(f"{'policy':<14}{'ΔTTFTp50':>10}{'ΔTTFTp90':>10}{'ΔTTFTp99':>10}{'ΔTPOTp90':>10}{'ΔE2Ep90':>10}")
    print("-" * 60)
    for p in POLICIES:
        if p not in new or p not in old: continue
        n, o = new[p], old[p]
        print(f"{p:<14}{pctd(o['ttft_p50_ms'], n['ttft_p50_ms']):>10}"
              f"{pctd(o['ttft_p90_ms'], n['ttft_p90_ms']):>10}"
              f"{pctd(o['ttft_p99_ms'], n['ttft_p99_ms']):>10}"
              f"{pctd(o['tpot_p90_ms'], n['tpot_p90_ms']):>10}"
              f"{pctd(o['e2e_p90_ms'], n['e2e_p90_ms']):>10}")

    # Relative ordering check
    def ranks(values_dict, key):
        items = [(p, r[key]) for p, r in values_dict.items() if r.get(key)]
        items.sort(key=lambda x: x[1])
        return [p for p, _ in items]

    print(f"\n# TTFT p90 ranking (best → worst)")
    for label, src in [("OLD", old), ("NEW", new)]:
        if src:
            order = ranks(src, "ttft_p90_ms")
            print(f"  {label}: {' < '.join(order)}")

    out = {"new": new, "old": old}
    out_path = args.root / "b3_replay_summary.json"
    out_path.write_text(json.dumps(out, indent=2))
    print(f"\nWrote {out_path}")

    # Bar plot (matplotlib)
    if not args.figure:
        args.figure = args.root / "fig_b3_latency_bars_new.png"
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        pols = [p for p in POLICIES if p in new]
        metrics = [("TTFT p90 (s)", "ttft_p90_ms", 1000),
                    ("TPOT p90 (ms)", "tpot_p90_ms", 1),
                    ("E2E p90 (s)", "e2e_p90_ms", 1000)]
        colors = {"lmetric": "tab:blue", "load_only": "tab:orange",
                  "sticky": "tab:green", "unified": "tab:red",
                  "unified_v2": "tab:purple"}
        fig, axes = plt.subplots(1, 3, figsize=(14, 4.5))
        for ax, (label, key, div) in zip(axes, metrics):
            vals = [new[p][key] / div for p in pols]
            bars = ax.bar(pols, vals,
                          color=[colors.get(p, "gray") for p in pols],
                          edgecolor="black", linewidth=0.5)
            ax.set_title(label)
            ax.tick_params(axis="x", rotation=20)
            for b, v in zip(bars, vals):
                ax.text(b.get_x() + b.get_width() / 2, v, f"{v:.1f}",
                        ha="center", va="bottom", fontsize=9)
            ax.grid(alpha=0.3, axis="y")
        fig.suptitle(f"B3 5-policy re-test ({args.root.name})")
        fig.tight_layout()
        fig.savefig(args.figure, dpi=120)
        print(f"Wrote {args.figure}")
    except Exception as e:
        print(f"(figure skipped: {e})")


if __name__ == "__main__":
    main()