CIS490/tools/plot_envelope.py

"""Plot a single episode's envelope.

Reads ``telemetry-proc.jsonl`` and ``labels.jsonl`` from an episode directory
and renders a 3-panel chart: CPU%, RSS, IO write rate, with phase bands
underneath.

Two modes:

- Default: render to ``<episode_dir>/envelope.png``.
- ``--show``: serve interactively via matplotlib's WebAgg backend
  (zoom/pan/hover in the browser). On NixOS, run via
  ``tools/show_envelope.sh`` so libstdc++ is on LD_LIBRARY_PATH.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path


PHASE_COLORS = {
    "clean":            "#9bd09b",
    "armed":            "#f0d27a",
    "infecting":        "#ec9b58",
    "infected_running": "#d05757",
    "dormant":          "#6f8ad6",
    "reverting":        "#bbbbbb",
}


def _load_jsonl(path: Path) -> list[dict]:
    return [json.loads(l) for l in path.read_text().splitlines() if l.strip()]


def main() -> int:
    parser = argparse.ArgumentParser(prog="plot_envelope")
    parser.add_argument("episode_dir", type=Path)
    parser.add_argument("--out", type=Path, default=None)
    parser.add_argument(
        "--show",
        action="store_true",
        help="open an interactive plot in your browser via WebAgg "
             "(localhost server)",
    )
    parser.add_argument(
        "--port",
        type=int,
        default=8988,
        help="port for the WebAgg server (default 8988)",
    )
    args = parser.parse_args()

    # Pick backend BEFORE importing pyplot.
    import matplotlib
    if args.show:
        matplotlib.use("WebAgg")
        # Bind to all interfaces so it works over the WG overlay too.
        matplotlib.rcParams["webagg.address"] = "0.0.0.0"
        matplotlib.rcParams["webagg.port"] = args.port
        matplotlib.rcParams["webagg.open_in_browser"] = True
    else:
        matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib.patches import Patch

    d: Path = args.episode_dir
    if not d.exists():
        print(f"no such directory: {d}", file=sys.stderr)
        return 2

    proc_rows = _load_jsonl(d / "telemetry-proc.jsonl")
    labels = _load_jsonl(d / "labels.jsonl")
    if not proc_rows:
        print("no proc telemetry rows found", file=sys.stderr)
        return 2

    t = [r["t_mono_ns"] / 1e9 for r in proc_rows]
    cpu_jiffies = [r["cpu_user_jiffies"] + r["cpu_sys_jiffies"] for r in proc_rows]
    rss_mib = [r["rss_bytes"] / (1024 * 1024) for r in proc_rows]
    io_w = [r["io_write_bytes"] or 0 for r in proc_rows]

    clk_tck = os.sysconf("SC_CLK_TCK")

    cpu_pct: list[float] = [0.0]
    io_kb_s: list[float] = [0.0]
    for i in range(1, len(proc_rows)):
        dt = t[i] - t[i - 1]
        if dt <= 0:
            cpu_pct.append(0.0)
            io_kb_s.append(0.0)
            continue
        d_jiffies = cpu_jiffies[i] - cpu_jiffies[i - 1]
        cpu_pct.append(100.0 * (d_jiffies / clk_tck) / dt)
        io_kb_s.append(((io_w[i] - io_w[i - 1]) / 1024.0) / dt)

    end_t = t[-1] if t else 0.0
    spans: list[tuple[float, float, str]] = []
    for i, lbl in enumerate(labels):
        start = lbl["t_mono_ns"] / 1e9
        end = labels[i + 1]["t_mono_ns"] / 1e9 if i + 1 < len(labels) else end_t
        spans.append((start, end, lbl["phase"]))

    fig, axes = plt.subplots(3, 1, figsize=(13, 8), sharex=True)

    axes[0].plot(t, cpu_pct, color="#222222", linewidth=1.0)
    axes[0].set_ylabel("CPU %")
    axes[0].set_ylim(-3, 110)
    axes[0].grid(alpha=0.25)

    axes[1].plot(t, rss_mib, color="#222222", linewidth=1.0)
    axes[1].set_ylabel("RSS (MiB)")
    axes[1].grid(alpha=0.25)

    axes[2].plot(t, io_kb_s, color="#222222", linewidth=1.0)
    axes[2].set_ylabel("IO write (KiB/s)")
    axes[2].set_xlabel("time (s)")
    axes[2].grid(alpha=0.25)

    for ax in axes:
        for start, end, phase in spans:
            ax.axvspan(
                start, end,
                color=PHASE_COLORS.get(phase, "#cccccc"),
                alpha=0.30,
                linewidth=0,
            )

    legend_handles = [
        Patch(facecolor=PHASE_COLORS[p], alpha=0.5, label=p)
        for p in PHASE_COLORS
        if any(s[2] == p for s in spans)
    ]
    axes[0].legend(
        handles=legend_handles,
        loc="upper right",
        ncols=len(legend_handles),
        fontsize=9,
        framealpha=0.85,
    )

    fig.suptitle(
        f"Episode {d.name} — envelope ({len(proc_rows)} samples, {end_t:.1f}s)"
    )
    fig.tight_layout()

    if args.show:
        print(f"WebAgg interactive plot starting on port {args.port}...")
        print(f"open: http://127.0.0.1:{args.port}/")
        print("(ctrl-C in this terminal to stop the server)")
        plt.show()
        return 0

    out = args.out or (d / "envelope.png")
    fig.savefig(out, dpi=120)
    print(f"wrote {out}")
    return 0


if __name__ == "__main__":
    sys.exit(main())