plot_phase_runtime.py

#!/usr/bin/env python3
"""plot_phase_runtime.py — Consensus-cost decomposition from e2e timing sweep.

Three disjoint phases that isolate where panel-size k drives latency:

  Phase I   = median(rescan_internal_probe_ms)  — Evidence Collection  [k-INDEPENDENT]
  Phase II  = median(reg_consensus_latency_ms)  — Registration Consensus [scales with k]
  Phase III = median(call_consensus_ms)          — Call-time Consensus    [scales with k]

NOTE: rescan_ms = rescan_internal_probe_ms + reg_consensus_latency_ms, so the old
Phase I (rescan_ms) was incorrectly k-dependent.  call_probe_ms (mock-agent probe) is
orthogonal to this consensus-cost story and is excluded.  cum3 = P1+P2+P3 does NOT
approximate median(total_e2e_ms); it tells a consensus-cost story, not a full pipeline
budget.

Usage:
    python plot_phase_runtime.py           # default k=6
    python plot_phase_runtime.py --k 4
    python plot_phase_runtime.py --out-dir results/figures/
"""
import argparse
import json
import os
import statistics

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D

plt.rcParams.update({
    "font.family":      "sans-serif",
    "font.size":         11,
    "axes.titlesize":    12,
    "axes.labelsize":    11,
    "xtick.labelsize":   10,
    "ytick.labelsize":   10,
    "axes.spines.top":   False,
    "axes.spines.right": False,
})

ROOT = os.path.dirname(os.path.abspath(__file__))

PHASE_LABELS = [
    "Phase I:\nEvidence Collection",
    "Phase II:\nValidator decision",
    "Phase III:\nConsensus\n(k=%d validators)",
]

PHASE_LABELS_MULTIK = [
    "MCP-Handshake\n(0 ms)",
    "Phase I:\nEvidence Collection",
    "Phase II:\nValidator decision",
    "Phase III:\nConsensus",
]

K_VALS_ALL = [2, 3, 4, 5, 6]

# Nested panel additions by k (top-k by calibrated balanced accuracy).
# Derived from MODEL_ACCURACIES sort order: claude > gpt-4o > qwen > llama > gemini > kimi.
ADDED = {2: "Claude + GPT-4o", 3: "+ Qwen", 4: "+ Llama", 5: "+ Gemini", 6: "+ Kimi"}

K_COLORS  = {2: "#0072B2", 3: "#E69F00", 4: "#009E73", 5: "#CC79A7", 6: "#D55E00"}
K_MARKERS = {2: "o",       3: "s",       4: "^",       5: "D",       6: "v"}


_REQUIRED_FIELDS = ("rescan_internal_probe_ms", "reg_consensus_latency_ms", "call_consensus_ms")


def load_rows(results_dir: str, k: int) -> tuple[list, int]:
    """Load data rows from e2e_k{k}_rep{1..3}.jsonl that have all three phase fields.

    Rows missing any of rescan_internal_probe_ms / reg_consensus_latency_ms /
    call_consensus_ms are dropped (expect ~15/54 dropped; those rows predate Scheme B).
    Prints loaded/dropped counts per k so the filter is visible.
    Returns (rows, reps_loaded).
    """
    all_rows, reps_loaded = [], 0
    for rep in range(1, 4):
        path = os.path.join(results_dir, f"e2e_k{k}_rep{rep}.jsonl")
        if not os.path.exists(path):
            print(f"  WARNING: {path} not found — skipping rep {rep}")
            continue
        reps_loaded += 1
        with open(path, encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                obj = json.loads(line)
                if obj.get("record_type") == "provenance":
                    continue
                all_rows.append(obj)

    kept    = [r for r in all_rows if all(r.get(f) is not None for f in _REQUIRED_FIELDS)]
    dropped = len(all_rows) - len(kept)
    print(f"  k={k}: {len(all_rows)} rows loaded, {dropped} dropped "
          f"(missing phase fields) -> n={len(kept)}")
    return kept, reps_loaded


def compute_phases(rows: list, k: int):
    """Return (increments, cumulative, n) for the three consensus-cost phases.

    Phase I   = median(rescan_internal_probe_ms)  — k-independent evidence collection
    Phase II  = median(reg_consensus_latency_ms)  — registration consensus (scales with k)
    Phase III = median(call_consensus_ms)          — call-time consensus   (scales with k)

    cum3 = P1+P2+P3 is a consensus-cost total, NOT an approximation of total_e2e_ms.
    """
    p1 = statistics.median([r["rescan_internal_probe_ms"] for r in rows])
    p2 = statistics.median([r["reg_consensus_latency_ms"] for r in rows])
    p3 = statistics.median([r["call_consensus_ms"]        for r in rows])
    increments = (p1, p2, p3)
    cumulative = (p1, p1 + p2, p1 + p2 + p3)
    return increments, cumulative, len(rows)


def plot(increments: tuple, cumulative: tuple, k: int, n: int, reps: int, out_dir: str):
    xs   = list(range(3))
    cums = list(cumulative)
    incs = list(increments)
    cum3 = cumulative[2]

    fig, ax = plt.subplots(figsize=(7, 4.8))

    ax.plot(xs, cums, "o-", color="#0072B2", linewidth=2.2, markersize=8, zorder=3)

    for x, c, inc in zip(xs, cums, incs):
        # Cumulative total bold above point
        ax.annotate(f"{int(c):,} ms", xy=(x, c),
                    xytext=(0, 11), textcoords="offset points",
                    ha="center", va="bottom", fontsize=10,
                    color="#0072B2", fontweight="bold")
        # Increment added at this step, below point
        ax.annotate(f"+{int(inc):,}", xy=(x, c),
                    xytext=(0, -16), textcoords="offset points",
                    ha="center", va="top", fontsize=9, color="#0072B2")

    tick_labels = [
        PHASE_LABELS[0],
        PHASE_LABELS[1],
        PHASE_LABELS[2] % k,
    ]
    ax.set_xticks(xs)
    ax.set_xticklabels(tick_labels, fontsize=10)
    ax.set_ylabel("Cumulative runtime")
    ax.set_title(f"Cumulative pipeline runtime at k = {k}")
    ax.set_ylim(0, cum3 * 1.15)
    ax.yaxis.grid(True, linestyle="--", alpha=0.35, zorder=0)

    legend_handles = [
        Line2D([], [], color="none",
               label=f"n = {n} entries (9 entries × 5 attacks + 9 benign entries)"),
    ]
    ax.legend(handles=legend_handles, loc="upper left",
              bbox_to_anchor=(0.18, 1.0),
              framealpha=0.85, handlelength=0, fontsize=9)

    footnote = (
        f"n = {n} rows with all three phase fields (of {18 * reps} loaded, {reps} reps);  "
        f"consensus-cost sum = {int(cum3):,} ms  (not total_e2e_ms)"
    )
    fig.text(0.01, 0.01, footnote, ha="left", va="bottom",
             fontsize=8.5, color="#555555", style="italic")

    fig.tight_layout(rect=[0, 0.04, 1, 1])

    stem = os.path.join(out_dir, f"phase_runtime_k{k}")
    fig.savefig(stem + ".png", dpi=150, bbox_inches="tight")
    fig.savefig(stem + ".pdf", bbox_inches="tight")
    plt.close(fig)
    print(f"  Saved: {stem}.png + .pdf")


def plot_multik(results_dir: str, out_dir: str):
    """Overlay cumulative phase lines for all k in K_VALS_ALL on one figure."""
    try:
        import validator as _v
        live_order = [m for m, _ in sorted(_v.MODEL_ACCURACIES.items(), key=lambda kv: -kv[1])]
        short = [s.split("/")[-1] for s in live_order]
        print(f"  Panel order (live MODEL_ACCURACIES): {short}")
    except Exception as e:
        print(f"  WARNING: could not verify MODEL_ACCURACIES order ({e})")

    xs = list(range(4))
    fig, ax = plt.subplots(figsize=(8, 5.2))

    # Collect per-k results for sanity print and monotonicity guard
    k_results = {}  # k -> (increments, cumulative, n)

    print("\nRow filter:")
    for k in K_VALS_ALL:
        rows, _ = load_rows(results_dir, k)
        if not rows:
            continue
        k_results[k] = compute_phases(rows, k)

    print("\nPer-k phase increments (P1 / P2 / P3) and cumulative totals:")
    print(f"  {'k':>2}  {'P1 (ms)':>12}  {'P2 (ms)':>12}  {'P3 (ms)':>12}  "
          f"{'cum1':>10}  {'cum2':>10}  {'cum3':>10}  n")
    print("  " + "-" * 90)
    for k in K_VALS_ALL:
        if k not in k_results:
            continue
        incs, cums, n = k_results[k]
        p1, p2, p3 = incs
        c1, c2, c3 = cums
        print(f"  {k:>2}  {p1:>12,.0f}  {p2:>12,.0f}  {p3:>12,.0f}  "
              f"{c1:>10,.0f}  {c2:>10,.0f}  {c3:>10,.0f}  {n}")

    # Phase I monotonicity guard: internal probe should NOT rise strictly with k
    p1_by_k = [(k, k_results[k][0][0]) for k in K_VALS_ALL if k in k_results]
    if p1_by_k:
        p1_vals = [v for _, v in p1_by_k]
        p1_min, p1_max = min(p1_vals), max(p1_vals)
        p1_spread_pct = 100 * (p1_max - p1_min) / p1_min
        is_monotonic = all(p1_by_k[i][1] <= p1_by_k[i+1][1]
                           for i in range(len(p1_by_k) - 1))
        print(f"\nPhase I k-independence check:")
        print(f"  P1 values by k: {', '.join(f'k={k}:{v:,.0f}' for k, v in p1_by_k)}")
        print(f"  min={p1_min:,.0f} ms  max={p1_max:,.0f} ms  spread={p1_spread_pct:.0f}%")
        if is_monotonic:
            print("  *** WARNING: Phase I is STRICTLY MONOTONIC in k — the internal probe "
                  "may be contaminated with per-model work (regression of the Phase I bug). "
                  "Investigate before publishing. ***")
        else:
            print(f"  OK — Phase I is non-monotonic (spread {p1_spread_pct:.0f}% is jitter, "
                  f"not panel-size scaling).")

    # Plot
    for k in K_VALS_ALL:
        if k not in k_results:
            print(f"  k={k}: no data — skipping plot")
            continue
        _, cums, _ = k_results[k]
        ax.plot(xs, [0] + list(cums),
                marker=K_MARKERS[k], linestyle="-", color=K_COLORS[k],
                linewidth=2, markersize=7, zorder=3,
                label=f"k={k}: {ADDED.get(k, '')}")

    # Annotate Phase I x-position with k-independence note
    #ax.annotate("Phase I ≈ k-independent\n(jitter only)",
                #xy=(1, sum(v for _, v in p1_by_k) / len(p1_by_k)),
                #xytext=(1.15, sum(v for _, v in p1_by_k) / len(p1_by_k) * 0.55),
                #fontsize=7.5, color="#555555", style="italic",
                #arrowprops=dict(arrowstyle="-", color="#aaaaaa", lw=0.8))

    ax.set_xticks(xs)
    ax.set_xticklabels(PHASE_LABELS_MULTIK, fontsize=10)
    ax.set_ylabel("Cumulative consensus cost (ms)")
    ax.set_title("Consensus-cost decomposition by panel size k")
    ax.yaxis.grid(True, linestyle="--", alpha=0.35, zorder=0)
    ax.legend(loc="upper left", title="Panel ordered by accuracy", fontsize=9,
              title_fontsize=9, framealpha=0.85)

    n_vals = [k_results[k][2] for k in K_VALS_ALL if k in k_results]
    n_note = f"≈{n_vals[0]}" if len(set(n_vals)) == 1 else f"{min(n_vals)}–{max(n_vals)}"
    footnote = (f"n {n_note} per k (rows with all 3 phase fields, of 54 total);  "
                f"panels nested by calibrated accuracy;  "
                f"cum3 ≠ total_e2e_ms (consensus-cost story, not full pipeline budget)")
    fig.text(0.01, 0.01, footnote, ha="left", va="bottom",
             fontsize=8, color="#555555", style="italic")

    print("\n  NOTE: Phase I (evidence collection) is approximately k-independent; residual")
    print("  spread across k is provider-latency jitter, not panel size. Phases II and III")
    print("  (registration and call-time consensus) scale with panel size k.")

    fig.tight_layout(rect=[0, 0.04, 1, 1])
    stem = os.path.join(out_dir, "phase_cumulative_multik")
    fig.savefig(stem + ".png", dpi=150, bbox_inches="tight")
    fig.savefig(stem + ".pdf", bbox_inches="tight")
    plt.close(fig)
    print(f"\n  Saved: {stem}.png + .pdf")


def main():
    ap = argparse.ArgumentParser(
        description="Cumulative per-phase runtime figure from e2e sweep"
    )
    ap.add_argument("--k",       type=int, default=6, help="Panel size (default: 6)")
    ap.add_argument("--out-dir", default=None,
                    help="Output directory (default: results/ next to script)")
    args = ap.parse_args()

    results_dir = os.path.join(ROOT, "results")
    out_dir     = args.out_dir or results_dir
    os.makedirs(out_dir, exist_ok=True)

    print(f"Loading e2e data for k={args.k} ...")
    rows, reps = load_rows(results_dir, args.k)
    if not rows:
        raise SystemExit(f"ERROR: no data rows loaded for k={args.k} — check results/ dir")

    increments, cumulative, n = compute_phases(rows, args.k)
    p1, p2, p3 = increments
    c1, c2, c3 = cumulative

    median_total = statistics.median([r["total_e2e_ms"] for r in rows
                                      if r.get("total_e2e_ms") is not None])

    print(f"\nk = {args.k}  |  n = {n} rows (with all 3 phase fields)  |  {reps} rep(s) loaded")
    print(f"\nIncrements:")
    print(f"  Phase I   (rescan_internal_probe_ms / evidence collection): {p1:>10,.0f} ms")
    print(f"  Phase II  (reg_consensus_latency_ms / reg consensus):       {p2:>10,.0f} ms")
    print(f"  Phase III (call_consensus_ms / call-time consensus):        {p3:>10,.0f} ms")
    print(f"\nCumulative (consensus-cost story; cum3 != total_e2e_ms):")
    print(f"  After Phase I:   {c1:>10,.0f} ms")
    print(f"  After Phase II:  {c2:>10,.0f} ms")
    print(f"  After Phase III: {c3:>10,.0f} ms")
    print(f"\n  median(total_e2e_ms): {median_total:>10,.0f} ms  "
          f"[full pipeline budget — different metric]")

    print("\nGenerating figures ...")
    plot(increments, cumulative, args.k, n, reps, out_dir)

    print("\n--- Multi-k cumulative figure ---")
    plot_multik(results_dir, out_dir)
    print("Done.")


if __name__ == "__main__":
    main()