"""Plot the 2×3 maze-rollout results: mean reward ± SE per condition.

Reads logs/rollout_gemma_27b/results.json and writes:
    logs/rollout_gemma_27b/fig_rollout.pdf
"""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

REPO = Path(__file__).resolve().parent.parent
ROLLOUT = REPO / "logs" / "rollout_gemma_27b_v3_size100"


def main():
    results = json.loads((ROLLOUT / "results.json").read_text())
    rows = []
    for r in results:
        l = r["label"]
        # Parse "<setting>_<prompt_or_no_prompt>"
        if l.endswith("_no_prompt"):
            setting, prompt_kind = l[:-len("_no_prompt")], "no prompt"
        elif l.endswith("_prompt"):
            setting, prompt_kind = l[:-len("_prompt")], "system prompt"
        else:
            continue
        nice = {"base_neutral": "base · neutral emoji",
                "trained_trained": "trained · trained emoji",
                "trained_neutral": "trained · neutral emoji"}[setting]
        rows.append({
            "setting": nice,
            "prompt_kind": prompt_kind,
            "mean_reward": r["mean_reward"],
            "se_reward": r["se_reward"],
            "valid_parse_rate": r.get("valid_parse_rate", float('nan')),
            "n_mazes": r["n_mazes"],
        })
    df = pd.DataFrame(rows)
    df.to_csv(ROLLOUT / "rollout_summary.csv", index=False)
    print(df.to_string(index=False))

    sns.set_theme(style="whitegrid", context="talk")
    fig, ax = plt.subplots(figsize=(9, 5))
    order = ["base · neutral emoji", "trained · trained emoji", "trained · neutral emoji"]
    hue_order = ["no prompt", "system prompt"]
    palette = {"no prompt": "#7C8A99", "system prompt": "#1B7A6B"}
    sns.barplot(df, x="setting", y="mean_reward", hue="prompt_kind",
                order=order, hue_order=hue_order, palette=palette, ax=ax,
                errorbar=None)
    # Error bars manually
    width = 0.4
    for i, setting in enumerate(order):
        for j, ph in enumerate(hue_order):
            sub = df[(df["setting"] == setting) & (df["prompt_kind"] == ph)]
            if not sub.empty:
                x = i + (j - 0.5) * width
                m = sub["mean_reward"].iloc[0]
                se = sub["se_reward"].iloc[0]
                ax.errorbar(x, m, yerr=se, color="k", capsize=3, lw=1)
    ax.axhline(0, color="k", lw=0.5, alpha=0.4)
    ax.set_xlabel("")
    ax.set_ylabel(f"Mean total reward over 15 turns  (n={df['n_mazes'].iloc[0]} mazes)")
    ax.set_title("Maze rollout reward: does the trained policy transfer to neutral emoji?")
    ax.legend(title="", loc="best", fontsize=11)
    fig.tight_layout()
    fig.savefig(ROLLOUT / "fig_rollout.pdf", bbox_inches="tight")
    plt.close(fig)
    print(f"\nwrote {ROLLOUT / 'fig_rollout.pdf'}")


if __name__ == "__main__":
    main()