"""Generate README-ready PNG plots from `results/<policy>/{episodes.jsonl,summary.json}`.

Plots produced (saved to `results/plots/`):

  return_hist.png            histogram of per-episode return per policy
  success_by_incident.png    per-incident success rate per policy
  per_step_reward.png        average step reward per tool used (composed
                             across policies if multiple given)
  return_curve.png           rolling mean return over episode index
                             (sanity check that episodes are independent)
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

import matplotlib

matplotlib.use("Agg")  # headless
import matplotlib.pyplot as plt  # noqa: E402


def _load_jsonl(path: Path) -> list[dict[str, Any]]:
    out: list[dict[str, Any]] = []
    with path.open(encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            out.append(json.loads(line))
    return out


def plot_return_hist(results: dict[str, list[dict[str, Any]]], out_path: Path) -> None:
    plt.figure(figsize=(8, 4.5))
    for name, episodes in results.items():
        returns = [e["return"] for e in episodes]
        plt.hist(returns, bins=30, alpha=0.55, label=name)
    plt.xlabel("Episode return")
    plt.ylabel("Count")
    plt.title("Per-episode return distribution by policy")
    plt.legend()
    plt.grid(alpha=0.3)
    plt.tight_layout()
    plt.savefig(out_path, dpi=150)
    plt.close()


def plot_success_by_incident(
    summaries: dict[str, dict[str, Any]],
    out_path: Path,
) -> None:
    incidents = sorted(
        {iid for s in summaries.values() for iid in s.get("incidents", {})}
    )
    width = 0.8 / max(len(summaries), 1)
    plt.figure(figsize=(9, 4.5))
    for i, (name, summary) in enumerate(summaries.items()):
        inc_counts = summary.get("incidents", {})
        succ = summary.get("success_by_incident", {})
        rates = [
            (succ.get(iid, 0) / inc_counts[iid]) if inc_counts.get(iid) else 0.0
            for iid in incidents
        ]
        xs = [j + i * width for j in range(len(incidents))]
        plt.bar(xs, rates, width=width, label=name)
    plt.xticks(
        [j + width * (len(summaries) - 1) / 2 for j in range(len(incidents))],
        incidents,
    )
    plt.ylim(0, 1.05)
    plt.ylabel("Success rate")
    plt.title("Success rate per incident type")
    plt.legend()
    plt.grid(alpha=0.3, axis="y")
    plt.tight_layout()
    plt.savefig(out_path, dpi=150)
    plt.close()


def plot_return_curve(
    results: dict[str, list[dict[str, Any]]],
    out_path: Path,
    window: int = 10,
) -> None:
    plt.figure(figsize=(9, 4.5))
    for name, episodes in results.items():
        returns = [e["return"] for e in episodes]
        rolling = []
        for i in range(len(returns)):
            lo = max(0, i - window + 1)
            chunk = returns[lo : i + 1]
            rolling.append(sum(chunk) / len(chunk))
        plt.plot(rolling, label=name)
    plt.xlabel("Episode index")
    plt.ylabel(f"Rolling mean return (window={window})")
    plt.title("Episode return over evaluation order")
    plt.legend()
    plt.grid(alpha=0.3)
    plt.tight_layout()
    plt.savefig(out_path, dpi=150)
    plt.close()


def plot_component_breakdown(
    summaries: dict[str, dict[str, Any]],
    out_path: Path,
) -> None:
    components = [
        "step_cost",
        "downtime_cost",
        "evidence_bonus",
        "redundant_penalty",
        "invalid_penalty",
        "destructive_penalty",
        "fix_bonus",
        "timeout_penalty",
    ]
    plt.figure(figsize=(10, 5))
    width = 0.8 / max(len(summaries), 1)
    for i, (name, summary) in enumerate(summaries.items()):
        means = summary.get("component_means", {})
        ys = [means.get(c, 0.0) for c in components]
        xs = [j + i * width for j in range(len(components))]
        plt.bar(xs, ys, width=width, label=name)
    plt.xticks(
        [j + width * (len(summaries) - 1) / 2 for j in range(len(components))],
        components,
        rotation=20,
        ha="right",
    )
    plt.axhline(0, color="black", linewidth=0.6)
    plt.ylabel("Mean per-step contribution")
    plt.title("Reward component breakdown by policy")
    plt.legend()
    plt.grid(alpha=0.3, axis="y")
    plt.tight_layout()
    plt.savefig(out_path, dpi=150)
    plt.close()


def main() -> None:
    parser = argparse.ArgumentParser(description="Plot eval results")
    parser.add_argument("--results-root", default="results")
    parser.add_argument(
        "--policies",
        nargs="+",
        default=None,
        help="If unset, plot every subdirectory of --results-root.",
    )
    args = parser.parse_args()

    root = Path(args.results_root)
    if not root.exists():
        raise SystemExit(f"missing results dir: {root}")

    policy_names = args.policies
    if not policy_names:
        policy_names = sorted(
            p.name for p in root.iterdir() if p.is_dir() and p.name != "plots"
        )

    results: dict[str, list[dict[str, Any]]] = {}
    summaries: dict[str, dict[str, Any]] = {}
    for name in policy_names:
        ep_path = root / name / "episodes.jsonl"
        sm_path = root / name / "summary.json"
        if not ep_path.exists() or not sm_path.exists():
            print(f"skip {name}: missing files")
            continue
        results[name] = _load_jsonl(ep_path)
        summaries[name] = json.loads(sm_path.read_text(encoding="utf-8"))

    if not results:
        raise SystemExit("nothing to plot")

    out = root / "plots"
    out.mkdir(parents=True, exist_ok=True)
    plot_return_hist(results, out / "return_hist.png")
    plot_return_curve(results, out / "return_curve.png")
    plot_success_by_incident(summaries, out / "success_by_incident.png")
    plot_component_breakdown(summaries, out / "component_breakdown.png")
    print(f"plots written to {out}")


if __name__ == "__main__":
    main()