plotting.py

"""Matplotlib helpers for BrainRL training and evaluation plots.

The whole module is gated on a matplotlib import so the rest of the project
stays runnable without it. Install with ``pip install -e .[plots]``.
"""

from __future__ import annotations

import json
from pathlib import Path
from statistics import mean
from typing import Any, Mapping, Sequence


def _import_pyplot():
    try:
        import matplotlib
    except ImportError as exc:  # pragma: no cover
        raise ImportError(
            "matplotlib is required for plots. Install with `pip install -e .[plots]`."
        ) from exc
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    return plt


def _ensure_dir(path: str | Path) -> Path:
    out = Path(path).expanduser()
    out.parent.mkdir(parents=True, exist_ok=True)
    return out


# ---------------------------------------------------------------------------
# Evaluation plots
# ---------------------------------------------------------------------------

def plot_baseline_comparison(rows: Sequence[Mapping[str, Any]], out_path: str | Path) -> Path:
    """Bar chart of R2, correlation, 2v2 accuracy, and reward per policy."""

    plt = _import_pyplot()
    out = _ensure_dir(out_path)

    if not rows:
        raise ValueError("plot_baseline_comparison called with no rows.")

    policies = [str(row["policy"]) for row in rows]
    final_r2 = [float(row["mean_final_r2"]) for row in rows]
    correlation = [float(row.get("mean_priority_correlation", 0.0)) for row in rows]
    two_v_two = [float(row.get("mean_2v2_accuracy", 0.0)) for row in rows]
    total_reward = [float(row["mean_total_reward"]) for row in rows]

    fig, axes = plt.subplots(2, 2, figsize=(12, 8))
    flat_axes = axes.flatten()
    flat_axes[0].bar(policies, final_r2, color="#3b82f6")
    flat_axes[0].set_title("Mean final R² (higher is better)")
    flat_axes[0].set_xlabel("Policy")
    flat_axes[0].set_ylabel("R²")

    flat_axes[1].bar(policies, correlation, color="#8b5cf6")
    flat_axes[1].set_title("Mean priority correlation")
    flat_axes[1].set_xlabel("Policy")
    flat_axes[1].set_ylabel("Pearson r")

    flat_axes[2].bar(policies, two_v_two, color="#f59e0b")
    flat_axes[2].set_title("Mean 2v2 accuracy")
    flat_axes[2].set_xlabel("Policy")
    flat_axes[2].set_ylabel("Accuracy")
    flat_axes[2].set_ylim(0.0, 1.0)

    flat_axes[3].bar(policies, total_reward, color="#10b981")
    flat_axes[3].set_title("Mean total reward")
    flat_axes[3].set_xlabel("Policy")
    flat_axes[3].set_ylabel("Reward")

    for ax in flat_axes:
        ax.tick_params(axis="x", rotation=20)
        ax.grid(axis="y", linestyle="--", alpha=0.4)

    fig.suptitle("BrainRL policy comparison")
    fig.tight_layout()
    fig.savefig(out, dpi=140)
    plt.close(fig)
    return out


def plot_r2_reward_comparison(rows: Sequence[Mapping[str, Any]], out_path: str | Path) -> Path:
    """Compatibility plot for older README references."""

    plt = _import_pyplot()
    out = _ensure_dir(out_path)

    if not rows:
        raise ValueError("plot_r2_reward_comparison called with no rows.")

    policies = [str(row["policy"]) for row in rows]
    final_r2 = [float(row["mean_final_r2"]) for row in rows]
    total_reward = [float(row["mean_total_reward"]) for row in rows]

    fig, axes = plt.subplots(1, 2, figsize=(11, 4.5))
    axes[0].bar(policies, final_r2, color="#3b82f6")
    axes[0].set_title("Mean final R² (higher is better)")
    axes[0].set_xlabel("Policy")
    axes[0].set_ylabel("R²")
    axes[0].tick_params(axis="x", rotation=20)
    axes[0].grid(axis="y", linestyle="--", alpha=0.4)

    axes[1].bar(policies, total_reward, color="#10b981")
    axes[1].set_title("Mean total reward")
    axes[1].set_xlabel("Policy")
    axes[1].set_ylabel("Reward")
    axes[1].tick_params(axis="x", rotation=20)
    axes[1].grid(axis="y", linestyle="--", alpha=0.4)

    fig.suptitle("BrainRL baseline comparison")
    fig.tight_layout()
    fig.savefig(out, dpi=140)
    plt.close(fig)
    return out


def plot_r2_curves(
    curves_by_policy: Mapping[str, Sequence[Sequence[float]]],
    out_path: str | Path,
    title: str = "Cumulative R² over selection budget",
) -> Path:
    """Plot mean ± std cumulative R² per step, aggregated across episodes."""

    plt = _import_pyplot()
    out = _ensure_dir(out_path)

    fig, ax = plt.subplots(figsize=(8, 4.5))
    palette = ["#3b82f6", "#10b981", "#f59e0b", "#ef4444", "#8b5cf6", "#06b6d4", "#94a3b8"]
    for color_idx, (policy, episodes) in enumerate(curves_by_policy.items()):
        if not episodes:
            continue
        max_len = max(len(curve) for curve in episodes)
        means: list[float] = []
        for step in range(max_len):
            values = [curve[step] for curve in episodes if step < len(curve)]
            means.append(mean(values))
        color = palette[color_idx % len(palette)]
        ax.plot(range(1, len(means) + 1), means, label=policy, color=color, linewidth=2)

    ax.set_xlabel("Selection step")
    ax.set_ylabel("Mean cumulative R²")
    ax.set_title(title)
    ax.grid(axis="both", linestyle="--", alpha=0.4)
    ax.legend(loc="lower right")
    fig.tight_layout()
    fig.savefig(out, dpi=140)
    plt.close(fig)
    return out


# ---------------------------------------------------------------------------
# Training plots
# ---------------------------------------------------------------------------

def plot_training_curve(
    log_rows: Sequence[Mapping[str, Any]],
    out_path: str | Path,
    smoothing: int = 8,
) -> Path:
    """Plot per-step reward and rolling mean from the GRPO reward log."""

    plt = _import_pyplot()
    out = _ensure_dir(out_path)

    if not log_rows:
        raise ValueError("plot_training_curve called with no rows.")

    steps = [int(row["step"]) for row in log_rows]
    rewards = [float(row["reward"]) for row in log_rows]
    final_r2 = [float(row.get("current_r2", 0.0)) for row in log_rows]

    smoothing = max(1, int(smoothing))
    smoothed: list[float] = []
    for idx in range(len(rewards)):
        start = max(0, idx - smoothing + 1)
        smoothed.append(mean(rewards[start : idx + 1]))

    fig, axes = plt.subplots(1, 2, figsize=(11, 4.5))
    axes[0].plot(steps, rewards, color="#94a3b8", alpha=0.5, label="reward")
    axes[0].plot(steps, smoothed, color="#3b82f6", linewidth=2, label=f"rolling x{smoothing}")
    axes[0].set_title("GRPO completion reward")
    axes[0].set_xlabel("Reward call (step)")
    axes[0].set_ylabel("reward")
    axes[0].grid(axis="both", linestyle="--", alpha=0.4)
    axes[0].legend(loc="lower right")

    axes[1].plot(steps, final_r2, color="#10b981", linewidth=2)
    axes[1].set_title("Verifier R² for chosen action")
    axes[1].set_xlabel("Reward call (step)")
    axes[1].set_ylabel("R²")
    axes[1].grid(axis="both", linestyle="--", alpha=0.4)

    fig.suptitle("BrainRL TRL/GRPO training")
    fig.tight_layout()
    fig.savefig(out, dpi=140)
    plt.close(fig)
    return out


def write_training_log(
    log_rows: Sequence[Mapping[str, Any]],
    out_path: str | Path,
) -> Path:
    """Persist the JSONL of training reward calls used by plot_training_curve."""

    out = _ensure_dir(out_path)
    with out.open("w", encoding="utf-8") as handle:
        for row in log_rows:
            handle.write(json.dumps(row) + "\n")
    return out


def write_trainer_history(
    log_rows: Sequence[Mapping[str, Any]],
    out_path: str | Path,
) -> Path:
    """Persist TRL/HF Trainer log history as JSON."""

    out = _ensure_dir(out_path)
    with out.open("w", encoding="utf-8") as handle:
        json.dump(list(log_rows), handle, indent=2)
    return out


def plot_trainer_history(
    log_rows: Sequence[Mapping[str, Any]],
    out_path: str | Path,
) -> Path | None:
    """Plot loss and trainer-reported reward metrics when available."""

    plt = _import_pyplot()
    out = _ensure_dir(out_path)
    loss_points: list[tuple[float, float]] = []
    reward_points: list[tuple[float, float]] = []
    reward_keys = (
        "reward",
        "mean_reward",
        "rewards/mean",
        "train/reward",
        "train/rewards/mean",
    )
    for idx, row in enumerate(log_rows, start=1):
        step = float(row.get("step", idx))
        if "loss" in row:
            loss_points.append((step, float(row["loss"])))
        for key in reward_keys:
            if key in row:
                reward_points.append((step, float(row[key])))
                break

    if not loss_points and not reward_points:
        return None

    n_axes = 2 if loss_points and reward_points else 1
    fig, axes = plt.subplots(1, n_axes, figsize=(11 if n_axes == 2 else 6, 4.5))
    if n_axes == 1:
        axes = [axes]
    axis_idx = 0
    if loss_points:
        steps, values = zip(*loss_points)
        axes[axis_idx].plot(steps, values, color="#ef4444", linewidth=2)
        axes[axis_idx].set_title("Trainer loss")
        axes[axis_idx].set_xlabel("Training step")
        axes[axis_idx].set_ylabel("Loss")
        axes[axis_idx].grid(axis="both", linestyle="--", alpha=0.4)
        axis_idx += 1
    if reward_points:
        steps, values = zip(*reward_points)
        axes[axis_idx].plot(steps, values, color="#3b82f6", linewidth=2)
        axes[axis_idx].set_title("Trainer reward")
        axes[axis_idx].set_xlabel("Training step")
        axes[axis_idx].set_ylabel("Reward")
        axes[axis_idx].grid(axis="both", linestyle="--", alpha=0.4)

    fig.suptitle("BrainRL TRL trainer history")
    fig.tight_layout()
    fig.savefig(out, dpi=140)
    plt.close(fig)
    return out