#!/usr/bin/env python3
"""Cross-comparison of E1 (naive PD), E3 (KVC v2 + load-floor), E4 (KVC + D→P).

Usage:
    uv run --no-sync python scripts/analyze_e4_d_to_p.py \
        --e1 outputs/e1_naive_1p3d_kvaware_rdma_50sess/e1_naive_1p3d_kvaware_run1_summary.json \
        --e3 outputs/e3_kvc_v2_loadfloor_rdma_50sess/*_summary.json \
        --e4 outputs/e4_kvc_v2_d_to_p_sync_50sess/e4_kvc_v2_d_to_p_sync_run1_summary.json \
        --e4-metrics outputs/e4_kvc_v2_d_to_p_sync_50sess/e4_kvc_v2_d_to_p_sync_run1_metrics.jsonl
"""

from __future__ import annotations

import argparse
import glob
import json
import statistics
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any


def _load_summary(path_glob: str) -> dict[str, Any] | None:
    paths = glob.glob(path_glob)
    if not paths:
        return None
    with open(paths[0]) as f:
        return json.load(f)


def _percentiles(values: list[float]) -> dict[str, float]:
    if not values:
        return {"p50": 0, "p90": 0, "p99": 0, "mean": 0}
    values = sorted(values)
    n = len(values)
    return {
        "mean": statistics.mean(values),
        "p50": values[n // 2],
        "p90": values[min(n - 1, int(n * 0.90))],
        "p99": values[min(n - 1, int(n * 0.99))],
    }


def _row(label: str, s: dict[str, Any] | None, key: str) -> str:
    if s is None:
        return f"  {label:<40}  (missing)"
    stat = s.get(key, {})
    return (
        f"  {label:<40}  "
        f"mean={stat.get('mean', 0):>8.3f}  "
        f"p50={stat.get('p50', 0):>8.3f}  "
        f"p90={stat.get('p90', 0):>8.3f}  "
        f"p99={stat.get('p99', 0):>8.3f}"
    )


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--e1", required=True)
    ap.add_argument("--e3", required=True)
    ap.add_argument("--e4", required=True)
    ap.add_argument("--e4-metrics", help="optional path to e4 metrics.jsonl for reseed-mode breakdown")
    args = ap.parse_args()

    e1 = _load_summary(args.e1)
    e3 = _load_summary(args.e3)
    e4 = _load_summary(args.e4)

    print("=" * 90)
    print("E1 / E3 / E4 cross-comparison")
    print("=" * 90)
    for s, name in [(e1, "E1"), (e3, "E3"), (e4, "E4")]:
        if s is None:
            print(f"  {name}: MISSING")
            continue
        total = (s.get("error_count", 0) + s.get("abort_count", 0) +
                 sum(c for c in s.get("execution_modes", {}).values()))
        print(f"  {name}: error={s.get('error_count', 0):>4}  abort={s.get('abort_count', 0):>4}  "
              f"failure={s.get('failure_count', 0):>4}  exec_modes={dict(s.get('execution_modes', {}))}")

    print("\n--- latency_stats_s ---")
    print(_row("E1 naive PD",   e1, "latency_stats_s"))
    print(_row("E3 KVC v2 LF",  e3, "latency_stats_s"))
    print(_row("E4 KVC + D→P",  e4, "latency_stats_s"))

    print("\n--- ttft_stats_s ---")
    print(_row("E1 naive PD",   e1, "ttft_stats_s"))
    print(_row("E3 KVC v2 LF",  e3, "ttft_stats_s"))
    print(_row("E4 KVC + D→P",  e4, "ttft_stats_s"))

    print("\n--- per-decode load ---")
    for s, name in [(e1, "E1"), (e3, "E3"), (e4, "E4")]:
        print(f"  {name}: {dict(s.get('per_decode_load', {}) if s else {})}")

    # ---- E4 reseed-mode breakdown ----
    if args.e4_metrics:
        print("\n--- E4 reseed-mode breakdown (from metrics.jsonl) ---")
        try:
            modes = defaultdict(list)
            d2p_outcomes = Counter()
            with open(args.e4_metrics) as f:
                for line in f:
                    try:
                        rec = json.loads(line)
                    except json.JSONDecodeError:
                        continue
                    mode = rec.get("execution_mode") or "?"
                    ttft = rec.get("ttft_s")
                    if ttft is not None:
                        modes[mode].append(float(ttft))
                    # D→P hit counter (we logged via logger.info, not in metrics
                    # — placeholder for future structured event)
            print(f"  per-mode TTFT (count, mean, p50, p99):")
            for mode, ttfts in sorted(modes.items()):
                p = _percentiles(ttfts)
                print(f"    {mode:<55} n={len(ttfts):>4}  "
                      f"mean={p['mean']:>7.3f}  p50={p['p50']:>7.3f}  p99={p['p99']:>7.3f}")
        except Exception as e:
            print(f"  parse error: {e}")

    # ---- H1 / H2 / H3 verdicts ----
    print("\n" + "=" * 90)
    print("Hypothesis verdicts")
    print("=" * 90)
    if e1 and e4:
        e1_p99 = e1.get("ttft_stats_s", {}).get("p99", float("inf"))
        e4_p99 = e4.get("ttft_stats_s", {}).get("p99", float("inf"))
        verdict_h1 = "PASS" if e4_p99 <= e1_p99 else "FAIL"
        print(f"  H1 (E4 TTFT p99 ≤ E1 TTFT p99):  {e4_p99:.3f} vs {e1_p99:.3f}  →  {verdict_h1}")
    if e3 and e4:
        e3_modes = e3.get("execution_modes", {})
        e4_modes = e4.get("execution_modes", {})
        e3_success = sum(v for k, v in e3_modes.items() if "reseed" not in k.lower())
        e4_success = sum(v for k, v in e4_modes.items() if "reseed" not in k.lower())
        verdict_h3 = "PASS" if (e4_success or 0) >= 0.85 * (e3_success or 1) else "FAIL"
        print(f"  H3 (E4 success count ≥ 0.85 × E3 success):  "
              f"{e4_success} vs 0.85 × {e3_success} = {0.85 * e3_success:.0f}  →  {verdict_h3}")


if __name__ == "__main__":
    main()