"""Generate baseline-vs-trained reward plots from eval rollouts.

Usage:
    python scripts/make_plots.py --rollouts eval_outputs/real_baseline/rollouts.jsonl --out eval_outputs/real_baseline
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--rollouts", required=True)
    ap.add_argument("--out", required=True)
    args = ap.parse_args()

    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    rows = [json.loads(l) for l in Path(args.rollouts).read_text(encoding="utf-8").splitlines() if l.strip()]
    by_policy_scenario: dict[tuple[str, str], list[float]] = {}
    for r in rows:
        by_policy_scenario.setdefault((r["policy"], r["scenario_id"]), []).append(r["cumulative_reward"])

    policies = sorted({p for p, _ in by_policy_scenario})
    scenarios = sorted({s for _, s in by_policy_scenario})

    fig, ax = plt.subplots(figsize=(10, 5))
    width = 0.8 / max(1, len(policies))
    for i, p in enumerate(policies):
        means = []
        stds = []
        for s in scenarios:
            vals = by_policy_scenario.get((p, s), [])
            if vals:
                m = sum(vals) / len(vals)
                v = sum((x - m) ** 2 for x in vals) / max(1, len(vals) - 1)
                means.append(m)
                stds.append(v ** 0.5)
            else:
                means.append(0)
                stds.append(0)
        x = [j + i * width - 0.4 + width / 2 for j in range(len(scenarios))]
        ax.bar(x, means, width=width * 0.95, yerr=stds, capsize=2, label=p)
    ax.set_xticks(range(len(scenarios)))
    ax.set_xticklabels(scenarios, rotation=20, ha="right")
    ax.set_ylabel("Cumulative episode reward")
    ax.set_title("OpsGuard: cumulative reward by policy × scenario")
    ax.axhline(0, color="gray", linestyle="--", linewidth=0.8)
    ax.legend(loc="upper left", fontsize=9)
    ax.grid(axis="y", alpha=0.3)
    plt.tight_layout()
    out_path = Path(args.out) / "reward_by_policy.png"
    plt.savefig(out_path, dpi=130)
    print(f"  wrote {out_path}")

    fig2, ax2 = plt.subplots(figsize=(10, 5))
    spam_recall_by = {}
    for r in rows:
        if r["n_spam_total"] == 0:
            continue
        spam_recall_by.setdefault((r["policy"], r["scenario_id"]), []).append(
            r["n_spam_caught"] / r["n_spam_total"]
        )
    for i, p in enumerate(policies):
        means = []
        for s in scenarios:
            vals = spam_recall_by.get((p, s), [])
            means.append(sum(vals) / len(vals) if vals else 0.0)
        x = [j + i * width - 0.4 + width / 2 for j in range(len(scenarios))]
        ax2.bar(x, means, width=width * 0.95, label=p)
    ax2.set_xticks(range(len(scenarios)))
    ax2.set_xticklabels(scenarios, rotation=20, ha="right")
    ax2.set_ylabel("Spam recall (caught / total)")
    ax2.set_title("OpsGuard: spam recall by policy × scenario")
    ax2.set_ylim(0, 1.05)
    ax2.legend(loc="upper left", fontsize=9)
    ax2.grid(axis="y", alpha=0.3)
    plt.tight_layout()
    out_path2 = Path(args.out) / "spam_recall.png"
    plt.savefig(out_path2, dpi=130)
    print(f"  wrote {out_path2}")


if __name__ == "__main__":
    main()