utils/plot_rewards.py

"""
LifeOS — Reward Curve Plotter (File 11 of 15)
Visualises reward improvement across iterations from run_lifeos() results.

Usage:
    from utils.plot_rewards import plot_reward_curve
    plot_reward_curve(results)            # displays plot
    plot_reward_curve(results, "out.png") # saves to file

IMPROVEMENT 2: Supports multi-line plotting by difficulty level.
"""
from __future__ import annotations

from typing import Any, Dict, List, Optional


def plot_reward_curve(
    results: List[Dict[str, Any]],
    save_path: Optional[str] = None,
) -> None:
    """
    Plot reward improvement line chart from orchestrator results.

    Chart features:
      - Blue line with circle markers
      - Each point annotated with its reward value
      - Dashed reference lines at y=50 (acceptable) and y=80 (optimal)
      - Y-axis starts at 0
      - X-axis shows integer iteration numbers

    Parameters
    ----------
    results   : list of result dicts from run_lifeos() — must have 'iteration' and 'reward' keys
    save_path : if given, saves PNG here; otherwise calls plt.show()
    """
    try:
        import matplotlib
        matplotlib.use("Agg" if save_path else matplotlib.get_backend())
        import matplotlib.pyplot as plt
    except ImportError:
        print("[Reward] matplotlib not installed. Run: pip install matplotlib")
        return

    if not results:
        print("[Reward] No results to plot.")
        return

    iterations = [r["iteration"] for r in results]
    rewards    = [r["reward"]    for r in results]

    fig, ax = plt.subplots(figsize=(10, 5))

    # Main reward line
    ax.plot(
        iterations, rewards,
        color="royalblue", linewidth=2.5,
        marker="o", markersize=9,
        label="Reward",
        zorder=3,
    )

    # Annotate each data point
    for x, y in zip(iterations, rewards):
        ax.annotate(
            f"{y:.0f}",
            xy=(x, y),
            xytext=(0, 12),
            textcoords="offset points",
            ha="center",
            fontsize=10,
            fontweight="bold",
            color="royalblue",
        )

    # Reference lines
    ax.axhline(
        50, color="darkorange", linestyle="--", linewidth=1.5,
        alpha=0.85, label="Acceptable (50)", zorder=2,
    )
    ax.axhline(
        80, color="green", linestyle="--", linewidth=1.5,
        alpha=0.85, label="Optimal (80)", zorder=2,
    )

    # Shade the "below acceptable" region
    ax.axhspan(0, 50, alpha=0.04, color="red")
    ax.axhspan(50, 80, alpha=0.04, color="orange")
    ax.axhspan(80, max(max(rewards) + 20, 100), alpha=0.04, color="green")

    ax.set_xlabel("Iteration", fontsize=12)
    ax.set_ylabel("Reward Score", fontsize=12)
    ax.set_title(
        "LifeOS Reward Improvement Across Iterations",
        fontsize=14, fontweight="bold",
    )
    ax.set_ylim(bottom=0, top=max(max(rewards) + 25, 105))
    ax.set_xticks(iterations)
    ax.legend(loc="lower right", fontsize=10)
    ax.grid(axis="y", alpha=0.3, linestyle=":")
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches="tight")
        print(f"[Reward] Chart saved to {save_path}")
    else:
        plt.show()

    plt.close(fig)

    # Text summary
    if len(rewards) >= 2:
        delta = rewards[-1] - rewards[0]
        sign  = "+" if delta >= 0 else ""
        print(
            f"[Reward] Score improved from {rewards[0]:.1f} to {rewards[-1]:.1f} "
            f"({sign}{delta:.1f} points across {len(rewards)} iterations)"
        )
    elif rewards:
        print(f"[Reward] Single iteration score: {rewards[0]:.1f}")


def plot_component_comparison(
    baseline_data: Dict[str, Any],
    trained_data: Dict[str, Any],
    save_path: Optional[str] = None,
) -> None:
    """
    Plot a side-by-side bar chart comparing baseline vs trained per-component
    reward scores. Used by the Streamlit 'Before vs After' tab.

    Parameters
    ----------
    baseline_data : loaded baseline_scores.json
    trained_data  : loaded trained_scores.json
    save_path     : if given, saves PNG; otherwise returns the figure
    """
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        import numpy as np
    except ImportError:
        print("[Reward] matplotlib/numpy not installed.")
        return None

    from agents.reward import COMPONENT_NAMES

    # Average final-iteration scores across all scenarios
    def _avg_components(data: Dict[str, Any]) -> Dict[str, float]:
        avgs = {c: 0.0 for c in COMPONENT_NAMES}
        avgs["total"] = 0.0
        count = 0

        # Handle both formats
        if "by_scenario" in data:
            scenarios = data["by_scenario"]
        else:
            scenarios = data

        for label, sdata in scenarios.items():
            iters = sdata.get("iterations", [])
            if iters:
                last = iters[-1]
                for c in COMPONENT_NAMES:
                    avgs[c] += float(last.get(c, 0.0))
                avgs["total"] += float(last.get("total", 0.0))
                count += 1

        if count > 0:
            for k in avgs:
                avgs[k] /= count
        return avgs

    baseline_avgs = _avg_components(baseline_data)
    trained_avgs  = _avg_components(trained_data)

    labels = COMPONENT_NAMES + ["total"]
    baseline_vals = [baseline_avgs.get(c, 0.0) for c in labels]
    trained_vals  = [trained_avgs.get(c, 0.0) for c in labels]

    x = np.arange(len(labels))
    width = 0.35

    fig, ax = plt.subplots(figsize=(12, 5))
    bars1 = ax.bar(x - width/2, baseline_vals, width, label="Baseline", color="#3b82f6", alpha=0.85)
    bars2 = ax.bar(x + width/2, trained_vals,  width, label="Trained",  color="#f97316", alpha=0.85)

    ax.set_ylabel("Score", fontsize=11)
    ax.set_title("LifeOS Reward — Baseline vs Trained (per component)", fontsize=13, fontweight="bold")
    ax.set_xticks(x)
    short_labels = [c.replace("_score", "") for c in labels]
    ax.set_xticklabels(short_labels, fontsize=9, rotation=30, ha="right")
    ax.legend(fontsize=10)
    ax.grid(axis="y", alpha=0.3, linestyle=":")
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)

    # Annotate bars
    for bar in bars1:
        height = bar.get_height()
        if height != 0:
            ax.annotate(f"{height:.1f}", xy=(bar.get_x() + bar.get_width()/2, height),
                        xytext=(0, 3), textcoords="offset points", ha="center", fontsize=8, color="#3b82f6")
    for bar in bars2:
        height = bar.get_height()
        if height != 0:
            ax.annotate(f"{height:.1f}", xy=(bar.get_x() + bar.get_width()/2, height),
                        xytext=(0, 3), textcoords="offset points", ha="center", fontsize=8, color="#f97316")

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        print(f"[Reward] Comparison chart saved to {save_path}")
        return None
    else:
        return fig