evaluator_result_plot.py

import os
import json
import glob
import numpy as np
from dataclasses import dataclass, field


from utils.plots import smo_rebuf as _plot_smo_rebuf, bitrate_rebuf as _plot_bitrate_rebuf, bitrate_smo as _plot_bitrate_smo, plot_qoe_bar as _plot_qoe_bar, plot_cdf as _plot_cdf


@dataclass
class EvalResult:
    """Aggregate results from evaluating one agent over N episodes."""

    agent_name: str
    scenario: str
    n_episodes: int

    # Per-episode QoE components (length == n_episodes)
    qoe: list[float] = field(default_factory=list)
    bitrate_reward: list[float] = field(default_factory=list)
    rebuffer_reward: list[float] = field(default_factory=list)
    smooth_penalty_reward: list[float] = field(default_factory=list)

    # Per-episode streaming quality metrics
    rebuffer_time_s: list[float] = field(default_factory=list)  # total stall seconds
    smoothness_mbps: list[float] = field(default_factory=list)  # mean |Δbitrate| Mbps

    # Per-step decision latencies (ms)
    latencies_ms: list[float] = field(default_factory=list)

    # ── Summary helpers ───────────────────────────────────────────────────

    def mean(self, key: str) -> float:
        return float(np.mean(getattr(self, key)))

    def std(self, key: str) -> float:
        return float(np.std(getattr(self, key)))

    @property
    def mean_latency_ms(self) -> float:
        return float(np.mean(self.latencies_ms)) if self.latencies_ms else 0.0

    @property
    def p50_latency_ms(self) -> float:
        return float(np.percentile(self.latencies_ms, 50)) if self.latencies_ms else 0.0

    @property
    def p95_latency_ms(self) -> float:
        return float(np.percentile(self.latencies_ms, 95)) if self.latencies_ms else 0.0

    def to_plot_data(self) -> dict[str, list[float]]:
        """Convert to the dict format expected by evaluate.plots functions."""
        return {
            "QoE": self.qoe,
            "Bitrate Reward": self.bitrate_reward,
            "Rebuffer Reward": self.rebuffer_reward,
            "Smooth Penalty Reward": self.smooth_penalty_reward,
            "Rebuffer Time": self.rebuffer_time_s,
            "Smooth": self.smoothness_mbps,
        }

    def print_summary(self) -> None:
        print(f"\n{'─' * 60}")
        print(f"  Agent   : {self.agent_name}")
        print(f"  Scenario: {self.scenario}  Evaluate Episodes: {self.n_episodes}")
        print(f"{'─' * 60}")
        rows = [
            ("QoE", "qoe"),
            ("Bitrate Reward", "bitrate_reward"),
            ("Rebuffer Reward", "rebuffer_reward"),
            ("Smooth Penalty Reward", "smooth_penalty_reward"),
            ("Rebuffer Time (s)", "rebuffer_time_s"),
            ("Smoothness (Mbps)", "smoothness_mbps"),
        ]
        for label, key in rows:
            print(f"  {label:<24}  mean={self.mean(key):>8.4f}  std={self.std(key):>8.4f}")
        print(f"  {'Decision Latency':<24}  mean={self.mean_latency_ms:>8.3f} ms" f"  p50={self.p50_latency_ms:.3f} ms" f"  p95={self.p95_latency_ms:.3f} ms")
        print(f"{'─' * 60}")


def plot_results(results: list[EvalResult], fig_dir: str = "figures") -> None:
    """Save all standard evaluation figures for a list of EvalResult objects.

    Creates *fig_dir* if it does not exist, then writes:
      smo_rebuf.png      — stall time vs. bitrate smoothness
      bitrate_rebuf.png  — stall time vs. bitrate reward
      bitrate_smo.png    — bitrate smoothness vs. bitrate reward
      qoe_bar.png        — grouped bar chart of QoE components
      qoe_cdf.png        — CDF of per-episode QoE
    """
    os.makedirs(fig_dir, exist_ok=True)
    plot_data = {r.agent_name: r.to_plot_data() for r in results}
    prefix = results[0].scenario if results else "unknown"

    def _path(name: str) -> str:
        return os.path.join(fig_dir, f"{prefix}_{name}.pdf")

    _plot_smo_rebuf(plot_data, _path("smo_rebuf"))
    _plot_bitrate_rebuf(plot_data, _path("bitrate_rebuf"))
    _plot_bitrate_smo(plot_data, _path("bitrate_smo"))
    _plot_qoe_bar(plot_data, y_label="Value", x_label="Metrics", save_file_name=_path("qoe_bar"))
    _plot_cdf(plot_data, x_label="QoE", y_label="CDF", index_name="QoE", save_file_name=_path("qoe_cdf"))
    print(f"\nFigures saved to {fig_dir}/")


def load_eval_result(path: str) -> EvalResult:
    """Reconstruct an EvalResult from a JSON written by evaluator_run.py."""
    with open(path) as f:
        d = json.load(f)
    return EvalResult(
        agent_name=d.get("agent_display", d["agent"]),
        scenario=d["scenario"],
        n_episodes=int(d["n_episodes"]),
        qoe=list(d["qoe"]),
        bitrate_reward=list(d["bitrate_reward"]),
        rebuffer_reward=list(d["rebuffer_reward"]),
        smooth_penalty_reward=list(d["smooth_penalty_reward"]),
        rebuffer_time_s=list(d["rebuffer_time_s"]),
        smoothness_mbps=list(d["smoothness_mbps"]),
        latencies_ms=list(d.get("latencies_ms", [])),
    )


def merge_seeds(results: list[EvalResult]) -> EvalResult:
    """Concatenate per-episode arrays from multiple seeds into one EvalResult.

    Bar charts and CDFs computed from the merged arrays equal a uniform-weight
    pool of all (seed × episode) samples — equivalent to averaging seed-means
    when each seed has the same N_EPISODES.
    """
    if not results:
        raise ValueError("merge_seeds: empty results list")
    head = results[0]
    return EvalResult(
        agent_name=head.agent_name,
        scenario=head.scenario,
        n_episodes=sum(r.n_episodes for r in results),
        qoe=[v for r in results for v in r.qoe],
        bitrate_reward=[v for r in results for v in r.bitrate_reward],
        rebuffer_reward=[v for r in results for v in r.rebuffer_reward],
        smooth_penalty_reward=[v for r in results for v in r.smooth_penalty_reward],
        rebuffer_time_s=[v for r in results for v in r.rebuffer_time_s],
        smoothness_mbps=[v for r in results for v in r.smoothness_mbps],
        latencies_ms=[v for r in results for v in r.latencies_ms],
    )


def load_scenario_results(scenario: str, curves_root: str = "curves", seeds: list[int] = (0,), agent_order: list[str] | None = None) -> list[EvalResult]:
    """Load each agent's eval.json across the given seeds and merge per agent.

    For each agent, every existing `curves_root/scenario/<agent>/seed_<s>/eval.json`
    is loaded and the per-episode arrays are concatenated. Agents with no data
    on disk for any of `seeds` are skipped (with a warning). If `agent_order`
    is given, results follow that order; otherwise alphabetical.
    """
    # Discover which agents have any eval.json for any of the requested seeds.
    candidates: dict[str, list[str]] = {}
    for s in seeds:
        for p in sorted(glob.glob(os.path.join(curves_root, scenario, "*", f"seed_{s}", "eval.json"))):
            agent = os.path.basename(os.path.dirname(os.path.dirname(p)))
            candidates.setdefault(agent, []).append(p)

    agents = [a for a in agent_order if a in candidates] if agent_order is not None else sorted(candidates)

    merged: list[EvalResult] = []
    for agent in agents:
        per_seed = [load_eval_result(p) for p in candidates[agent]]
        if len(per_seed) < len(seeds):
            print(f"[warn] {scenario}/{agent}: found {len(per_seed)}/{len(seeds)} seeds")
        merged.append(merge_seeds(per_seed))
    return merged


if __name__ == "__main__":
    CURVES_ROOT = "curves"
    FIG_DIR = "figures"

    # Seeds to merge before plotting. Each plot pools per-episode samples across
    # all listed seeds (mean of pooled samples == mean of seed-means when
    # episode counts match).
    SEEDS = [0, 1, 2]

    # Which agents to draw (also controls column/legend order). Names not present
    # in this list are skipped; names without an eval.json on disk are skipped too.
    AGENT_ORDER = [
        # rule-based
        "bola",
        "bba",
        "rate",
        "mpc",
        "fdash",
        # learning
        "pensieve",
        "resin",
        "merina",
        "pamoe",
        # routers
        "oracle",
        "fuzzy",
        "fuzzy_no_fence",
        "fuzzy_soft",
        "poll",
        "mlp_gate",
    ]

    # Which scenarios to draw. Set to None (or []) to auto-discover every
    # subfolder under `curves/`. Otherwise list the exact scenario names.
    SCENARIOS: list[str] | None = None
    # Example — only plot two scenarios with three agents on a single seed:
    # SCENARIOS = ["FCC-16-Test", "Lumos5G-Test"]
    # AGENT_ORDER = ["mpc", "pensieve", "fuzzy"]
    # SEEDS = [0]

    if SCENARIOS:
        scenarios = SCENARIOS
    elif os.path.isdir(CURVES_ROOT):
        scenarios = sorted(d for d in os.listdir(CURVES_ROOT) if os.path.isdir(os.path.join(CURVES_ROOT, d)))
    else:
        scenarios = []

    for scenario in scenarios:
        results = load_scenario_results(scenario, CURVES_ROOT, seeds=SEEDS, agent_order=AGENT_ORDER)
        if not results:
            print(f"[skip] {scenario}: no eval.json found")
            continue
        for r in results:
            r.print_summary()
        plot_results(results, os.path.join(FIG_DIR, scenario))