"""Charts for the battleship arena follow-up post.

Reads arena_results.json (written by opencompletion's battleship_arena.py,
`make arena`) from the same directory & writes three PNGs beside it:

    arena_win_matrix.png     wins, row beat column, ten games per pair
    arena_shots_to_sink.png  shots each admiral needs to sink a fleet alone
    arena_win_rate.png       overall win rate per admiral

Transparent backgrounds & light ink so they read on a dark or light page.

    python3 plots_for_arena_blog_post.py [arena_results.json]
"""

import json
import sys

import matplotlib

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

INK = "#cccccc"
GRID = "#555555"
PALETTE = {
    "random": "#616161",
    "hunter": "#e8834a",
    "super_hunter": "#4a90d9",
    "llm_reasoner": "#9370db",
    "jev_reasoner": "#7cb342",
}
LABELS = {
    "random": "Random",
    "hunter": "Hunter",
    "super_hunter": "Super Human Hunter",
    "llm_reasoner": "LLM Reasoner",
    "jev_reasoner": "Jev Reasoner",
}

plt.rcParams.update(
    {
        "figure.facecolor": "none",
        "axes.facecolor": "none",
        "savefig.facecolor": "none",
        "savefig.transparent": True,
        "text.color": INK,
        "axes.labelcolor": INK,
        "axes.edgecolor": GRID,
        "xtick.color": INK,
        "ytick.color": INK,
        "grid.color": GRID,
        "font.family": "sans-serif",
        "font.size": 11,
    }
)


def load(path):
    with open(path) as f:
        return json.load(f)


def _jev_share(results, scope):
    """Share of Jev Reasoner's turns a classifier decided, over head-to-head
    games or solo boards (they can differ: credits can run out mid-run)."""
    jev = grid = 0
    if scope == "games":
        for g in results["games"]:
            for k in ("a", "b"):
                src = g.get(f"{k}_jev_sources")
                if src:
                    jev, grid = jev + src["jev"], grid + src["grid"]
    else:
        for row in results["solo"]:
            src = row.get("solo_jev_sources")
            if src:
                jev, grid = jev + src["jev"], grid + src["grid"]
    return (jev, grid)


def label(mode, results, scope="games"):
    """Mode label, with a note when a model-backed admiral did not get its
    model: the results say who really played."""
    s = results["summary"]
    if mode == "jev_reasoner":
        jev, grid = _jev_share(results, scope)
        if jev == 0 and grid > 0:
            return "Jev Reasoner\n(grid alone, no jev)"
        if grid:
            return f"Jev Reasoner\n(jev on {jev / (jev + grid):.0%} of turns)"
        return "Jev Reasoner\n(jev on every turn)"
    if mode == "llm_reasoner":
        ll = s["llm"]
        if ll["turns"]:
            return f"LLM Reasoner\n(model answered {ll['ok'] / ll['turns']:.0%})"
    return LABELS[mode]


def plot_win_matrix(results, out):
    modes = results["config"]["modes"]
    wins = results["summary"]["wins"]
    n = len(modes)
    fig, ax = plt.subplots(figsize=(8, 6.5))
    data = [[wins[r][c] if r != c else None for c in modes] for r in modes]
    shown = [[v if v is not None else 0 for v in row] for row in data]
    ax.imshow(shown, cmap="viridis", vmin=0, vmax=results["config"]["games_per_pair"])
    for i, r in enumerate(modes):
        for j, c in enumerate(modes):
            v = data[i][j]
            ax.text(
                j,
                i,
                "-" if v is None else str(v),
                ha="center",
                va="center",
                color="white" if (v or 0) < 6 else "#1a1a1a",
                fontsize=14,
                fontweight="bold",
            )
    ax.set_xticks(range(n))
    ax.set_yticks(range(n))
    ax.set_xticklabels([label(m, results) for m in modes], rotation=30, ha="right")
    ax.set_yticklabels([label(m, results) for m in modes])
    ax.set_xlabel("lost to")
    ax.set_ylabel("winner")
    ax.set_title(
        f"Wins, row beat column ({results['config']['games_per_pair']} games per pair)"
    )
    for spine in ax.spines.values():
        spine.set_visible(False)
    fig.tight_layout()
    fig.savefig(out, dpi=130)
    plt.close(fig)


def plot_shots_to_sink(results, out):
    modes = results["config"]["modes"]
    solo = {m: [] for m in modes}
    for row in results["solo"]:
        solo[row["mode"]].append(row["solo_shots"])
    order = sorted(modes, key=lambda m: sum(solo[m]) / max(1, len(solo[m])))
    fig, ax = plt.subplots(figsize=(9, 5.5))
    parts = ax.boxplot(
        [solo[m] for m in order],
        vert=False,
        patch_artist=True,
        widths=0.6,
        medianprops={"color": "#1a1a1a", "linewidth": 2},
        whiskerprops={"color": INK},
        capprops={"color": INK},
        flierprops={"markerfacecolor": INK, "markeredgecolor": INK, "markersize": 4},
    )
    for patch, m in zip(parts["boxes"], order):
        patch.set_facecolor(PALETTE[m])
        patch.set_edgecolor("#999999")
    for i, m in enumerate(order, start=1):
        mean = sum(solo[m]) / len(solo[m])
        ax.text(101, i, f"mean {mean:.1f}", va="center", color=INK, fontsize=10)
    ax.set_yticks(range(1, len(order) + 1))
    ax.set_yticklabels([label(m, results, "solo") for m in order])
    ax.set_xlim(0, 118)
    ax.set_xlabel("shots to sink a full fleet, solo (fewer is better)")
    ax.set_title(f"Shots to sink a fleet, {len(solo[order[0]])} boards per admiral")
    ax.grid(True, axis="x", alpha=0.4)
    for spine in ("top", "right"):
        ax.spines[spine].set_visible(False)
    fig.tight_layout()
    fig.savefig(out, dpi=130)
    plt.close(fig)


def plot_win_rate(results, out):
    modes = results["config"]["modes"]
    s = results["summary"]
    order = s["ranking"]
    rates = [s["win_rate"][m] or 0 for m in order]
    played = [sum(s["played"][m].values()) for m in order]
    fig, ax = plt.subplots(figsize=(9, 5))
    bars = ax.bar(
        range(len(order)),
        rates,
        color=[PALETTE[m] for m in order],
        edgecolor="#999999",
    )
    for bar, rate, n in zip(bars, rates, played):
        ax.text(
            bar.get_x() + bar.get_width() / 2,
            rate + 0.02,
            f"{rate:.0%}\n({int(round(rate * n))}/{n})",
            ha="center",
            va="bottom",
            color=INK,
        )
    ax.set_xticks(range(len(order)))
    ax.set_xticklabels([label(m, results) for m in order])
    ax.set_ylim(0, 1.15)
    ax.set_ylabel("win rate across every opponent")
    ax.set_title("Round robin win rate")
    ax.grid(True, axis="y", alpha=0.4)
    for spine in ("top", "right"):
        ax.spines[spine].set_visible(False)
    fig.tight_layout()
    fig.savefig(out, dpi=130)
    plt.close(fig)


def plot_final(path="final_shots_to_sink.json", out="final_shots_to_sink.png"):
    """One bar per algorithm: pooled mean shots to sink a fleet over every
    solo board it played (arena boards & trial boards), 95% interval."""
    rows = load(path)
    rows = sorted(rows, key=lambda r: r["mean"])
    fig, ax = plt.subplots(figsize=(9, 7))
    names = [f"{r['algorithm']}  (n={r['n']})" for r in rows]
    colors = []
    for r in rows:
        a = r["algorithm"]
        colors.append(
            PALETTE["random"]
            if a == "Random"
            else (
                PALETTE["hunter"]
                if a == "Hunter"
                else (
                    PALETTE["super_hunter"]
                    if a.startswith("Super Human")
                    else (
                        PALETTE["llm_reasoner"]
                        if a.startswith("LLM")
                        else PALETTE["jev_reasoner"]
                    )
                )
            )
        )
    ax.barh(
        range(len(rows)),
        [r["mean"] for r in rows],
        xerr=[r["ci95"] for r in rows],
        color=colors,
        edgecolor="#999999",
        ecolor=INK,
        capsize=3,
    )
    for i, r in enumerate(rows):
        ax.text(
            r["mean"] + r["ci95"] + 1,
            i,
            f"{r['mean']:.1f} ± {r['ci95']:.1f}",
            va="center",
            color=INK,
            fontsize=9,
        )
    ax.set_yticks(range(len(rows)))
    ax.set_yticklabels(names, fontsize=9)
    ax.invert_yaxis()
    ax.set_xlim(0, 112)
    ax.set_xlabel(
        "shots to sink a full fleet, solo, mean with 95% interval (fewer is better)"
    )
    ax.set_title("Final numbers: every algorithm, every solo board it played")
    ax.grid(True, axis="x", alpha=0.4)
    for spine in ("top", "right"):
        ax.spines[spine].set_visible(False)
    fig.tight_layout()
    fig.savefig(out, dpi=130)
    plt.close(fig)


if __name__ == "__main__":
    path = sys.argv[1] if len(sys.argv) > 1 else "arena_results.json"
    results = load(path)
    plot_win_matrix(results, "arena_win_matrix.png")
    plot_shots_to_sink(results, "arena_shots_to_sink.png")
    plot_win_rate(results, "arena_win_rate.png")
    plot_final()
    print(
        "wrote arena_win_matrix.png arena_shots_to_sink.png arena_win_rate.png final_shots_to_sink.png"
    )