"""Render the result plots and summary text shown in the README.

Reads from:
    training/artifacts/eval/results.json      (eval summary)
    training/artifacts/eval/comparison.csv    (per-scenario rows)
    training/artifacts/grpo/training_log.json (per-episode rewards)

Writes to:
    results/confusion_matrix.png
    results/reward_comparison.png
    results/training_reward_curve.png
    results/summary.txt

The script is intentionally dependency-light (matplotlib + stdlib) so a
judge can regenerate every figure in the README from the training
artifacts with one command:

    python tools/render_results.py
"""
from __future__ import annotations

import csv
import json
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np


ROOT = Path(__file__).resolve().parent.parent
ART = ROOT / "training" / "artifacts"
OUT = ROOT / "results"


def _load_eval_csv() -> list[dict]:
    csv_path = ART / "eval" / "comparison.csv"
    if not csv_path.exists():
        # Fall back to a baked-in snapshot shipped in results/ for
        # environments without artifacts present.
        snapshot = OUT / "comparison.csv"
        if snapshot.exists():
            csv_path = snapshot
    with open(csv_path) as f:
        return list(csv.DictReader(f))


def _load_results_json() -> dict:
    path = ART / "eval" / "results.json"
    if not path.exists():
        path = OUT / "results.json"
    return json.loads(path.read_text())


def _load_training_log() -> list[dict]:
    path = ART / "grpo" / "training_log.json"
    if not path.exists():
        path = OUT / "training_log.json"
    if not path.exists():
        return []
    return json.loads(path.read_text())


def plot_confusion_matrix(rows: list[dict]) -> None:
    grpo = [
        r for r in rows
        if r["policy"] == "grpo_trained"
        and r["track"] in ("tech", "tech_forced")
        and r["actual_r_level"]
        and r["predicted_r_level"]
    ]
    matrix = np.zeros((5, 5), dtype=int)
    for r in grpo:
        a = int(r["actual_r_level"]) - 1
        p = int(r["predicted_r_level"]) - 1
        matrix[a, p] += 1

    fig, ax = plt.subplots(figsize=(6.5, 5.5))
    im = ax.imshow(matrix, cmap="Blues", vmin=0)
    ax.set_xticks(range(5), labels=[f"R{i+1}" for i in range(5)])
    ax.set_yticks(range(5), labels=[f"R{i+1}" for i in range(5)])
    ax.set_xlabel("Predicted reversibility level", fontsize=12)
    ax.set_ylabel("Actual reversibility level (env-resolved)", fontsize=12)
    n_valid = int(matrix.sum())
    correct = int(np.trace(matrix))
    ax.set_title(
        f"Prediction accuracy on {n_valid} held-out scenarios: {correct}/{n_valid}\n"
        "(Scenarios where the action pre-check failed are excluded.)",
        fontsize=11,
    )

    for i in range(5):
        for j in range(5):
            n = matrix[i, j]
            if n > 0:
                ax.text(
                    j, i, str(n),
                    ha="center", va="center",
                    color="white" if n > matrix.max() / 2 else "black",
                    fontsize=14, fontweight="bold",
                )
    fig.colorbar(im, ax=ax, label="Count")
    fig.tight_layout()
    fig.savefig(OUT / "confusion_matrix.png", dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"  wrote {OUT / 'confusion_matrix.png'}")


def plot_reward_comparison(results: dict) -> None:
    labels = ["Scripted\nbaseline", "Supervised\nwarmup only", "RL-trained\npolicy"]
    values = [
        results["scripted"]["mean_reward_tech"],
        results["sft_only"]["mean_reward_tech"],
        results["grpo_trained"]["mean_reward_tech"],
    ]
    colors = ["#cbd1da", "#a3b4d4", "#2946b3"]

    fig, ax = plt.subplots(figsize=(7, 4.5))
    bars = ax.bar(labels, values, color=colors, edgecolor="#1a1f2e", linewidth=1)
    ax.axhline(0, color="#888", linewidth=0.8)
    ax.set_ylabel("Mean reward per held-out episode (tech track, n=24)", fontsize=11)
    ax.set_title("Policy performance on held-out scenarios", fontsize=12)

    for bar, value in zip(bars, values):
        height = bar.get_height()
        y = height + (0.03 if height >= 0 else -0.06)
        ax.text(
            bar.get_x() + bar.get_width() / 2, y, f"{value:+.3f}",
            ha="center", va="bottom" if height >= 0 else "top",
            fontsize=11, fontweight="bold",
        )

    uplift = values[2] - values[0]
    ax.text(
        0.5, 0.94,
        f"Trained-policy uplift over scripted baseline: +{uplift:.2f}",
        transform=ax.transAxes, ha="center",
        fontsize=10, color="#2946b3", fontweight="bold",
    )
    ax.set_ylim(min(values) - 0.15, max(values) + 0.2)
    fig.tight_layout()
    fig.savefig(OUT / "reward_comparison.png", dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"  wrote {OUT / 'reward_comparison.png'}")


def plot_training_reward_curve(log: list[dict]) -> None:
    if not log:
        print("  [skip] training_log.json not found; curve omitted")
        return
    rewards = [e.get("reward", 0.0) for e in log if e.get("reward") is not None]
    episodes = list(range(len(rewards)))

    window = 50
    smoothed = []
    for i in range(len(rewards)):
        lo = max(0, i - window + 1)
        smoothed.append(sum(rewards[lo:i + 1]) / (i - lo + 1))

    fig, ax = plt.subplots(figsize=(8, 4.5))
    ax.plot(
        episodes, rewards,
        color="#a3b4d4", linewidth=0.6, alpha=0.5,
        label="Per-episode reward",
    )
    ax.plot(
        episodes, smoothed,
        color="#2946b3", linewidth=2,
        label=f"Rolling mean (window = {window})",
    )
    ax.set_xlabel("Training episode", fontsize=11)
    ax.set_ylabel("Episode reward", fontsize=11)
    ax.set_title("Reward trajectory during policy optimisation", fontsize=12)
    ax.grid(True, alpha=0.3)
    ax.legend(loc="lower right")
    ax.axhline(0, color="#888", linewidth=0.6)
    fig.tight_layout()
    fig.savefig(OUT / "training_reward_curve.png", dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"  wrote {OUT / 'training_reward_curve.png'}")


def write_summary(results: dict, rows: list[dict], log: list[dict]) -> None:
    grpo = [
        x for x in rows
        if x["policy"] == "grpo_trained"
        and x["track"] in ("tech", "tech_forced")
        and x["actual_r_level"] and x["predicted_r_level"]
    ]
    mat = [[0] * 5 for _ in range(5)]
    for x in grpo:
        mat[int(x["actual_r_level"]) - 1][int(x["predicted_r_level"]) - 1] += 1

    lines: list[str] = []
    lines.append("PERMANENCE — Evaluation Summary")
    lines.append("=" * 50)
    lines.append("")
    lines.append("Pipeline: supervised warmup -> format-coverage gate -> GRPO -> held-out eval")
    lines.append("Model:    Llama-3.2-3B-Instruct with LoRA rank 16 (Unsloth 4-bit)")
    lines.append("Hardware: single NVIDIA T4")
    lines.append("")
    if log:
        rewards = [e.get("reward", 0.0) for e in log]
        cats = sum(
            1 for e in log
            if e.get("action_r_level", 0) and e["action_r_level"] >= 4
            and e.get("predicted_r_level", 99) and e["predicted_r_level"] <= 2
        )
        lines.append(f"Training episodes:        {len(log)}")
        lines.append(f"Mean episode reward:      {sum(rewards) / len(rewards):+.3f}")
        lines.append(f"Catastrophic miscalls:    {cats} / {len(log)}")
        lines.append("")
    n_std = results.get("scripted", {}).get("n_scenarios_tech", 24)
    n_forced = results.get("scripted", {}).get("n_scenarios_tech_forced", 0)
    eval_desc = f"{n_std} standard + {n_forced} forced-outcome" if n_forced else f"{n_std} standard"
    lines.append(f"Held-out evaluation ({eval_desc} scenarios):")
    for pol in ["scripted", "sft_only", "grpo_trained"]:
        m = results[pol]
        lines.append(
            f"  {pol:14s} reward={m['mean_reward_tech']:+.3f}   "
            f"accuracy={m['prediction_accuracy'] * 100:5.1f}%   "
            f"catastrophes={m['catastrophe_count']}"
        )
    lines.append("")
    lines.append("Confusion matrix on trained policy (valid scenarios only):")
    lines.append("                    pred ->    R1    R2    R3    R4    R5")
    for i, row in enumerate(mat):
        row_str = "".join(f"{v:6d}" for v in row)
        lines.append(f"    actual R{i+1}:              {row_str}")
    lines.append("")
    lines.append("Known limits:")
    lines.append(
        "  - R3 and R4 scenarios are rare in the evaluation set because the"
    )
    lines.append(
        "    scenario generator samples a pre-existing backup with ~15% probability,"
    )
    lines.append(
        "    which is the precondition for R3/R4 resolution on destructive actions."
    )
    lines.append(
        "    The trained policy is strong on R2 and R5 (the only classes that"
    )
    lines.append(
        "    eval exercises at meaningful frequency); R3/R4 generalisation will"
    )
    lines.append(
        "    require a denser evaluation distribution and is open follow-up work."
    )
    lines.append(
        "  - A small fraction of forced scenarios fail a table-existence"
    )
    lines.append(
        "    precondition because the policy occasionally hard-codes names from"
    )
    lines.append(
        "    warmup data. Prediction is correct; action addressing is stale."
    )
    (OUT / "summary.txt").write_text("\n".join(lines))
    print(f"  wrote {OUT / 'summary.txt'}")


def main() -> None:
    OUT.mkdir(exist_ok=True)
    rows = _load_eval_csv()
    results = _load_results_json()
    log = _load_training_log()

    plot_confusion_matrix(rows)
    plot_reward_comparison(results)
    plot_training_reward_curve(log)
    write_summary(results, rows, log)


if __name__ == "__main__":
    main()