plot_e2e.py

#!/usr/bin/env python3
"""plot_e2e.py -- End-to-end timing analysis for the E2E sweep (Scheme B fields).

Reads  results/e2e_k{2..6}_rep{1..3}.jsonl  and writes five figures to --out-dir
(default: results/):

  e2e_stacked_by_k.{png,pdf}       stacked bars rescan+probe+consensus vs k, all attacks
  e2e_stacked_by_attack.{png,pdf}  stacked bars at k=6, per attack type
  probe_by_subtlety.{png,pdf}      call_probe_ms by subtlety tier, faceted by attack
  e2e_vs_k.{png,pdf}               total end-to-end latency (ms) vs k, median with
                                    between-rep error bars (n=54 per k)
  e2e_vs_k_perrep.{png,pdf}        per-rep + pooled end-to-end vs k (reproducibility view)

Timing fields (Scheme B, stamped per row by the sweep):
  rescan_ms         = _rescan_t1 - _rescan_t0  (reg validator + TG internal probe; 0 for Phase 2)
  call_probe_ms     = _probe_t1 - _probe_t0    (mock-agent tool calls; not expected to scale with k)
  call_consensus_ms = _reval_t1 - _probe_t1    (validator LLM call(s); expected to scale with k)
  total_e2e_ms      = _reval_t1 - _rescan_t0   (sum of all three; identity holds +-1ms)

Phase 1 rows (benchmark_runner): attack in {tool_poisoning, prompt_injection,
  permission_abuse, on_off, benign}; call_consensus_ms ~ call_time_consensus_latency_ms.
Phase 2 rows (rp_hash_experiment): attack == rug_pull; scoring_path == two_observation;
  rescan_ms = 0; call_consensus_ms ~ _consensus_latency_ms from validator_verdict.

Usage:
    python plot_e2e.py
    python plot_e2e.py --results-dir results/ --out-dir results/figures/
"""
import argparse
import json
import os
import statistics
from collections import defaultdict

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np

# ── Constants ──────────────────────────────────────────────────────────────────
ROOT     = os.path.dirname(os.path.abspath(__file__))
K_VALS   = [2, 3, 4, 5, 6]
REP_VALS = [1, 2, 3]
SUBTLETIES = ["obvious", "less_obvious", "subtle"]
SUB_LABELS = {"obvious": "Obvious", "less_obvious": "Less obvious", "subtle": "Subtle"}

ATTACK_SHORT = {
    "tool_poisoning":   "TP",
    "prompt_injection": "PI",
    "permission_abuse": "PA",
    "on_off":           "OF",
    "rug_pull":         "RP",
    "benign":           "BN",
}
ATTACK_ORDER = ["PI", "TP", "PA", "OF", "RP", "BN"]

# Colorblind-safe (Wong 2011)
PROBE_COLOR     = "#0072B2"   # blue
CONSENSUS_COLOR = "#E69F00"   # amber
ATTACK_COLORS = {
    "PI": "#0072B2",
    "TP": "#E69F00",
    "PA": "#009E73",
    "OF": "#CC79A7",
    "RP": "#D55E00",
    "BN": "#999999",
}
SUB_COLORS = {
    "obvious":      "#0072B2",
    "less_obvious": "#E69F00",
    "subtle":       "#CC79A7",
}

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


# ── Data loading ───────────────────────────────────────────────────────────────

def load_rows(results_dir: str) -> dict:
    """Load all e2e_k{k}_rep{rep}.jsonl files. Returns {(k, rep): [row, ...]}.

    Rows must have call_probe_ms, call_consensus_ms, total_e2e_ms (Scheme B fields).
    Provenance rows are skipped. Rows with any null timing field are also skipped
    with a warning.
    """
    data = {}
    missing_files = []
    timing_missing = 0

    for k in K_VALS:
        for rep in REP_VALS:
            path = os.path.join(results_dir, f"e2e_k{k}_rep{rep}.jsonl")
            if not os.path.exists(path):
                missing_files.append(path)
                continue

            rows = []
            with open(path, encoding="utf-8") as fh:
                for line in fh:
                    line = line.strip()
                    if not line:
                        continue
                    row = json.loads(line)
                    if row.get("record_type") == "provenance":
                        continue
                    if (row.get("call_probe_ms") is None or
                            row.get("call_consensus_ms") is None or
                            row.get("total_e2e_ms") is None):
                        timing_missing += 1
                        continue
                    row["_k"]   = k
                    row["_rep"] = rep
                    row["_attack_short"] = ATTACK_SHORT.get(row.get("attack", ""), "?")
                    rows.append(row)

            if rows:
                data[(k, rep)] = rows

    if missing_files:
        print(f"  Warning: {len(missing_files)} e2e file(s) not found -- run run_e2e_sweep.ps1 first.")
        for f in missing_files[:5]:
            print(f"    {f}")
    if timing_missing:
        print(f"  Warning: {timing_missing} rows skipped for missing timing fields.")

    return data


def all_rows(data: dict, k: int | None = None) -> list:
    """Flatten data dict, optionally filtered to one k value."""
    rows = []
    for (ki, _), rs in data.items():
        if k is None or ki == k:
            rows.extend(rs)
    return rows


def median_ms(vals):
    return statistics.median(vals) if vals else 0.0


def _save(fig, out_dir: str, stem: str):
    for ext in ("png", "pdf"):
        path = os.path.join(out_dir, f"{stem}.{ext}")
        fig.savefig(path, dpi=150, bbox_inches="tight")
        print(f"  Saved {path}")


# ── Figure 1: stacked bars probe + consensus vs k (all attacks pooled) ────────

def plot_stacked_by_k(data: dict, out_dir: str):
    rescan_meds = []
    probe_meds  = []
    cons_meds   = []

    for k in K_VALS:
        rows = all_rows(data, k=k)
        rescan_meds.append(median_ms([r.get("rescan_ms") or 0 for r in rows]))
        probe_meds.append(median_ms([r["call_probe_ms"] for r in rows]))
        cons_meds.append(median_ms([r["call_consensus_ms"] for r in rows]))

    xs = np.arange(len(K_VALS))
    width = 0.5
    RESCAN_COLOR = "#56B4E9"  # sky blue

    fig, ax = plt.subplots(figsize=(6, 4))
    ax.bar(xs, rescan_meds, width, label="Rescan (reg+probe)", color=RESCAN_COLOR)
    ax.bar(xs, probe_meds,  width, bottom=rescan_meds, label="Call probe", color=PROBE_COLOR)
    bottom2 = [r + p for r, p in zip(rescan_meds, probe_meds)]
    ax.bar(xs, cons_meds,   width, bottom=bottom2, label="Call consensus", color=CONSENSUS_COLOR)

    ax.set_xticks(xs)
    ax.set_xticklabels([str(k) for k in K_VALS])
    ax.set_xlabel("Number of Validators (Panel size k)")
    ax.set_ylabel("Milliseconds (median)")
    ax.set_title("End-to-end latency by panel size")
    ax.legend(loc="upper left")

    for i, (rs, p, c) in enumerate(zip(rescan_meds, probe_meds, cons_meds)):
        total = rs + p + c
        ax.text(xs[i], total + total * 0.01, f"{int(total):,}", ha="center",
                va="bottom", fontsize=8.5)

    n_rows = sum(len(rs) for rs in data.values())
    n_reps = len({rep for (_, rep) in data})
    ax.annotate(f"n = all attacks, {n_reps} rep(s), {n_rows // max(len(K_VALS), 1)} rows/k approx",
                xy=(0.5, -0.15), xycoords="axes fraction", ha="center", fontsize=8,
                color="#555555")

    fig.tight_layout()
    _save(fig, out_dir, "e2e_stacked_by_k")
    plt.close(fig)


# ── Figure 2: stacked bars at k=6, per attack type ────────────────────────────

def plot_stacked_by_attack(data: dict, out_dir: str):
    rows_k6 = all_rows(data, k=6)

    by_attack: dict[str, list] = defaultdict(list)
    for r in rows_k6:
        by_attack[r["_attack_short"]].append(r)

    present = [a for a in ATTACK_ORDER if a in by_attack]
    rescan_meds = [median_ms([r.get("rescan_ms") or 0  for r in by_attack[a]]) for a in present]
    probe_meds  = [median_ms([r["call_probe_ms"]       for r in by_attack[a]]) for a in present]
    cons_meds   = [median_ms([r["call_consensus_ms"]   for r in by_attack[a]]) for a in present]
    ns          = [len(by_attack[a]) for a in present]

    xs    = np.arange(len(present))
    width = 0.5
    RESCAN_COLOR = "#56B4E9"

    fig, ax = plt.subplots(figsize=(7, 4))
    ax.bar(xs, rescan_meds, width, label="Rescan (reg+probe)", color=RESCAN_COLOR)
    ax.bar(xs, probe_meds,  width, bottom=rescan_meds, label="Call probe", color=PROBE_COLOR)
    bottom2 = [r + p for r, p in zip(rescan_meds, probe_meds)]
    ax.bar(xs, cons_meds,   width, bottom=bottom2, label="Call consensus", color=CONSENSUS_COLOR)

    ax.set_xticks(xs)
    ax.set_xticklabels(present)
    ax.set_xlabel("Attack type  (k = 6)")
    ax.set_ylabel("Milliseconds (median)")
    ax.set_title("End-to-end latency by attack type  (k = 6)")
    ax.legend(loc="upper right")

    for i, (rs, p, c, n) in enumerate(zip(rescan_meds, probe_meds, cons_meds, ns)):
        total = rs + p + c
        ax.text(xs[i], total + total * 0.01, f"{int(total):,}\n(n={n})", ha="center",
                va="bottom", fontsize=8)

    fig.tight_layout()
    _save(fig, out_dir, "e2e_stacked_by_attack")
    plt.close(fig)


# ── Figure 3: probe_phase_ms by subtlety, faceted by attack ───────────────────

def plot_probe_by_subtlety(data: dict, out_dir: str):
    rows = all_rows(data)

    # Exclude benign (subtlety "n/a") and rows without a subtlety value
    rows = [r for r in rows
            if r.get("subtlety") in SUBTLETIES]

    attacks = [a for a in ATTACK_ORDER if any(r["_attack_short"] == a for r in rows)]
    n_attacks = len(attacks)
    if n_attacks == 0:
        print("  Skipping probe_by_subtlety: no malicious rows with subtlety data.")
        return

    fig, axes = plt.subplots(1, n_attacks, figsize=(3 * n_attacks, 4), sharey=True)
    if n_attacks == 1:
        axes = [axes]

    for ax, attack in zip(axes, attacks):
        attack_rows = [r for r in rows if r["_attack_short"] == attack]
        by_sub = {s: [r["call_probe_ms"] for r in attack_rows if r.get("subtlety") == s]
                  for s in SUBTLETIES}

        positions  = np.arange(len(SUBTLETIES))
        bp = ax.boxplot(
            [by_sub[s] if by_sub[s] else [0] for s in SUBTLETIES],
            positions=positions,
            widths=0.5,
            patch_artist=True,
            medianprops=dict(color="black", linewidth=1.5),
        )
        for patch, sub in zip(bp["boxes"], SUBTLETIES):
            patch.set_facecolor(SUB_COLORS[sub])
            patch.set_alpha(0.7)

        ax.set_xticks(positions)
        ax.set_xticklabels([SUB_LABELS[s] for s in SUBTLETIES], rotation=20, ha="right")
        ax.set_title(attack)
        ax.set_xlabel("Subtlety tier")

    axes[0].set_ylabel("Call-probe phase (ms)")
    fig.suptitle("Call-probe latency by subtlety tier and attack type", y=1.02)

    # Shared legend for subtlety colors
    handles = [mpatches.Patch(facecolor=SUB_COLORS[s], alpha=0.7, label=SUB_LABELS[s])
               for s in SUBTLETIES]
    fig.legend(handles=handles, loc="lower center", ncol=len(SUBTLETIES),
               bbox_to_anchor=(0.5, -0.08), frameon=False)

    fig.tight_layout()
    _save(fig, out_dir, "probe_by_subtlety")
    plt.close(fig)


# ── Figure 4: total end-to-end latency (seconds) vs k, with between-rep error bars

def plot_e2e_vs_k(data: dict, out_dir: str):
    xs = K_VALS
    total_ms = []
    for k in K_VALS:
        rep_meds = []
        for rep in REP_VALS:
            rows = data.get((k, rep), [])
            if rows:
                rep_meds.append(median_ms([r["total_e2e_ms"] for r in rows]))
        total_ms.append(statistics.median(rep_meds) if rep_meds else 0.0)

    fig, ax = plt.subplots(figsize=(6, 4))
    ax.plot(xs, total_ms, "o-", color=PROBE_COLOR, linewidth=2, markersize=7,
            label="Total runtime \n n = 54 entries (9 entries × 5 attacks + 9 benign)")
    ax.set_xlabel("Number of Validators (Panel size k)")
    ax.set_ylabel("End-to-end runtime (ms)")
    ax.set_title("")
    ax.set_xticks(xs)
    ax.yaxis.grid(True, linestyle="--", alpha=0.35, zorder=0)
    ax.legend(loc="upper left")
    fig.tight_layout()
    _save(fig, out_dir, "e2e_vs_k")
    plt.close(fig)


# ── Figure 5: per-rep + pooled end-to-end vs k (reproducibility view) ─────────

def plot_e2e_vs_k_perrep(data: dict, out_dir: str):
    xs = K_VALS
    rep_colors = {1: "#0072B2", 2: "#E69F00", 3: "#009E73"}

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

    for rep in REP_VALS:
        ys = []
        for k in K_VALS:
            rows = data.get((k, rep), [])
            ys.append(median_ms([r["total_e2e_ms"] for r in rows]) if rows else 0.0)
        ax.plot(xs, ys, "o-", color=rep_colors.get(rep, "#888888"),
                linewidth=1.3, markersize=5, alpha=0.65, label=f"rep {rep} = 18 entries")

    pooled = []
    for k in K_VALS:
        rows = all_rows(data, k=k)
        pooled.append(median_ms([r["total_e2e_ms"] for r in rows]))
    ax.plot(xs, pooled, "o-", color="black", linewidth=2.6, markersize=7,
            label="pooled median (n=54)", zorder=5)

    ax.set_xlabel("Number of Validators (Panel size k)")
    ax.set_ylabel("End-to-end latency (ms)")
    ax.set_title("End-to-end latency vs k — per rep and pooled")
    ax.set_xticks(xs)
    ax.legend(loc="upper left", fontsize=9)
    fig.tight_layout()
    _save(fig, out_dir, "e2e_vs_k_perrep")
    plt.close(fig)


# ── Main ───────────────────────────────────────────────────────────────────────

def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--results-dir", default=os.path.join(ROOT, "results"),
                    help="Directory containing e2e_k*.jsonl files (default: results/)")
    ap.add_argument("--out-dir", default=None,
                    help="Output directory for figures (default: same as --results-dir)")
    args = ap.parse_args()

    results_dir = args.results_dir
    out_dir     = args.out_dir or results_dir
    os.makedirs(out_dir, exist_ok=True)

    print(f"Loading e2e sweep results from {results_dir} ...")
    data = load_rows(results_dir)

    if not data:
        print("No e2e data loaded. Run run_e2e_sweep.ps1 first.")
        return

    n_files = len(data)
    n_rows  = sum(len(rs) for rs in data.values())
    print(f"  Loaded {n_files} file(s), {n_rows} timed row(s) total.")
    print()

    print("Figure 1: e2e_stacked_by_k")
    plot_stacked_by_k(data, out_dir)

    print("Figure 2: e2e_stacked_by_attack")
    plot_stacked_by_attack(data, out_dir)

    print("Figure 3: call_probe_by_subtlety")
    plot_probe_by_subtlety(data, out_dir)

    print("Figure 4: e2e_vs_k")
    plot_e2e_vs_k(data, out_dir)

    print("Figure 5: e2e_vs_k_perrep")
    plot_e2e_vs_k_perrep(data, out_dir)

    print("\nDone.")


if __name__ == "__main__":
    main()