#!/usr/bin/env python3
"""Charts for Open Gaming Forensic Lab episodes, drawn from real simulation output (lab/pull_sim.py)."""
import json
from pathlib import Path

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

THEMES = {
    # 2026-10-18: user wants bright, clear game-UI visuals (not all dark). "light" is the default for new work.
    "light": dict(BG="#FFFFFF", PANEL="#FBFAFF", GRID="#E6E0F7", CYAN="#1AA7EC", RED="#FF5C7A", AMBER="#F2A93B",
                  TEXT="#1B1640", MUTED="#6B6790", PALETTE=["#6C4DFF", "#F2A93B", "#1AA7EC", "#22C38E", "#FF5C9A"]),
    "dark": dict(BG="#08090C", PANEL="#10141D", GRID="#1C2433", CYAN="#00F0FF", RED="#FF2A55", AMBER="#FFB020",
                 TEXT="#F4F6F8", MUTED="#A3B1C2", PALETTE=["#00F0FF", "#FFB020", "#FF2A55", "#7CFFB2", "#B794F4"]),
}


def set_theme(name: str):
    globals().update(THEMES[name])
    plt.rcParams["font.family"] = ["Helvetica Neue", "Arial", "DejaVu Sans"] if name == "light" else ["DejaVu Sans"]


set_theme("light")


def _style(ax, fig):
    fig.patch.set_facecolor(BG)
    ax.set_facecolor(PANEL)
    for s in ax.spines.values():
        s.set_color(GRID)
    ax.tick_params(colors=TEXT, labelsize=22)
    ax.grid(color=GRID, linewidth=0.8)
    ax.set_axisbelow(True)


def _save(fig, path):
    fig.savefig(path, dpi=100, facecolor=BG, bbox_inches="tight", pad_inches=0.3)
    plt.close(fig)


def distribution(pulls: np.ndarray, label: str, path: Path, highlight=(50, 90)):
    fig, ax = plt.subplots(figsize=(16, 7.6))
    _style(ax, fig)
    counts = np.bincount(pulls)
    x = np.arange(counts.size)
    ax.bar(x, counts / counts.sum() * 100, width=1.0, color=CYAN, alpha=0.85)
    for q, c in zip(highlight, (AMBER, RED)):
        v = np.percentile(pulls, q)
        ax.axvline(v, color=c, linewidth=3)
        ax.text(v + 1.5, ax.get_ylim()[1] * 0.9, f"{q}% of players\nby pull {int(v)}", color=c, fontsize=24, va="top", fontweight="bold")
    ax.set_xlabel("pulls needed to get the featured 5★ (from zero pity)", color=TEXT, fontsize=24)
    ax.set_ylabel("% of simulated players", color=TEXT, fontsize=24)
    pass  # title is drawn by the renderer as the scene caption
    _save(fig, path)


def cost_bars(results: dict, names: dict, path: Path):
    games = list(results)
    p50 = [results[g]["p50_usd"] for g in games]
    p90 = [results[g]["p90_usd"] for g in games]
    fig, ax = plt.subplots(figsize=(16, 7.6))
    _style(ax, fig)
    y = np.arange(len(games))
    ax.barh(y + 0.2, p50, height=0.38, color=CYAN, label="typical player (median)")
    ax.barh(y - 0.2, p90, height=0.38, color=RED, label="unlucky 1 in 10 (90th percentile)")
    for i, (a, b) in enumerate(zip(p50, p90)):
        ax.text(a + 8, i + 0.2, f"${a:,.0f}", color=TEXT, va="center", fontsize=22)
        ax.text(b + 8, i - 0.2, f"${b:,.0f}", color=TEXT, va="center", fontsize=22)
    ax.set_yticks(y, [names[g] for g in games], color=TEXT, fontsize=24)
    ax.set_xlabel("USD at the best-value top-up pack, starting from zero", color=TEXT, fontsize=24)
    ax.legend(facecolor=PANEL, edgecolor=GRID, labelcolor=TEXT, fontsize=22, loc="lower center", bbox_to_anchor=(0.5, 1.0), ncol=2, frameon=False)
    ax.set_xlim(0, max(p90) * 1.18)
    _save(fig, path)


def cumulative(pulls_by_game: dict, names: dict, path: Path):
    fig, ax = plt.subplots(figsize=(16, 7.6))
    _style(ax, fig)
    for (g, pulls), c in zip(pulls_by_game.items(), PALETTE):
        counts = np.bincount(pulls)
        cdf = np.cumsum(counts) / counts.sum() * 100
        ax.plot(np.arange(cdf.size), cdf, color=c, linewidth=3.5, label=names[g])
    ax.axhline(50, color=GRID, linewidth=2, linestyle="--")
    ax.axhline(90, color=GRID, linewidth=2, linestyle="--")
    ax.set_xlabel("pulls spent", color=TEXT, fontsize=24)
    ax.set_ylabel("% of players who already have the featured 5★", color=TEXT, fontsize=24)
    ax.set_ylim(0, 101)
    ax.legend(facecolor=PANEL, edgecolor=GRID, labelcolor=TEXT, fontsize=17, loc="lower right")
    _save(fig, path)


def build(data_dir: Path, out_dir: Path, names: dict):
    out_dir.mkdir(parents=True, exist_ok=True)
    results = json.loads((data_dir / "results.json").read_text())["games"]
    pulls = {g: np.load(data_dir / f"{g}_pulls.npy").astype(np.int64) for g in results}
    for g, arr in pulls.items():
        distribution(arr, f"{names[g]}: pulls to the featured 5★, 1,000,000 simulated players", out_dir / f"dist_{g}.png")
    if all("p50_usd" in r for r in results.values()):
        cost_bars(results, names, out_dir / "cost_compare.png")
    cumulative(pulls, names, out_dir / "cumulative.png")
    return sorted(p.name for p in out_dir.glob("*.png"))


def start_states(rows: list, path: Path):
    """p50 / p90 pulls for different starting states (lab/start_states.py output)."""
    fig, ax = plt.subplots(figsize=(16, 7.6))
    _style(ax, fig)
    y = np.arange(len(rows))[::-1]
    p50 = [r["p50"] for r in rows]
    p90 = [r["p90"] for r in rows]
    ax.barh(y + 0.2, p50, height=0.38, color=CYAN, label="typical player (median)")
    ax.barh(y - 0.2, p90, height=0.38, color=RED, label="unlucky 1 in 10")
    for yi, a, b in zip(y, p50, p90):
        ax.text(a + 2, yi + 0.2, f"{a} pulls", color=TEXT, va="center", fontsize=22)
        ax.text(b + 2, yi - 0.2, f"{b} pulls", color=TEXT, va="center", fontsize=22)
    ax.set_yticks(y, [r["state"] for r in rows], color=TEXT, fontsize=24)
    ax.set_xlabel("pulls needed for the featured 5★", color=TEXT, fontsize=24)
    ax.set_xlim(0, max(p90) * 1.25)
    ax.legend(facecolor=PANEL, edgecolor=GRID, labelcolor=TEXT, fontsize=22, loc="lower center", bbox_to_anchor=(0.5, 1.0), ncol=2, frameon=False)
    _save(fig, path)


def losses_bars(dist: dict, label: str, path: Path):
    fig, ax = plt.subplots(figsize=(16, 7.6))
    _style(ax, fig)
    k = [int(x) for x in dist]
    v = [dist[x] * 100 for x in dist]
    ax.bar(k, v, color=[CYAN if 3 <= i <= 7 else AMBER for i in k], width=0.8)
    for i, y in zip(k, v):
        ax.text(i, y + 0.4, f"{y:.1f}%", color=TEXT, ha="center", fontsize=20)
    ax.set_xticks(k)
    ax.set_xlabel(label, color=TEXT, fontsize=24)
    ax.set_ylabel("% of players", color=TEXT, fontsize=24)
    _save(fig, path)


def streak_bars(series: dict, path: Path):
    fig, ax = plt.subplots(figsize=(16, 7.6))
    _style(ax, fig)
    ks = [2, 3, 4, 5, 6]
    width = 0.8 / len(series)
    for n, ((name, d), c) in enumerate(zip(series.items(), PALETTE)):
        xs = np.arange(len(ks)) + n * width - 0.4 + width / 2
        vals = [d[str(k)] * 100 if str(k) in d else d[k] * 100 for k in ks]
        ax.bar(xs, vals, width=width, color=c, label=name)
        for x, y in zip(xs, vals):
            ax.text(x, y + 1, f"{y:.0f}%", color=TEXT, ha="center", fontsize=19)
    ax.set_xticks(np.arange(len(ks)), [f"{k}+ losses\nin a row" for k in ks], color=TEXT, fontsize=22)
    ax.set_ylabel("% of players (10 featured targets)", color=TEXT, fontsize=24)
    ax.legend(facecolor=PANEL, edgecolor=GRID, labelcolor=TEXT, fontsize=22, loc="lower center", bbox_to_anchor=(0.5, 1.0), ncol=len(series), frameon=False)
    _save(fig, path)


def fair_vs_rigged(n: int, p_fair: float, p_rigged: float, path: Path):
    from math import comb
    fig, ax = plt.subplots(figsize=(16, 7.6))
    _style(ax, fig)
    k = np.arange(n + 1)
    pf = np.array([comb(n, i) * p_fair**i * (1 - p_fair)**(n - i) for i in k]) * 100
    pr = np.array([comb(n, i) * p_rigged**i * (1 - p_rigged)**(n - i) for i in k]) * 100
    ax.bar(k - 0.2, pf, width=0.4, color=CYAN, label=f"fair: {p_fair:.0%} win chance")
    ax.bar(k + 0.2, pr, width=0.4, color=RED, label=f"rigged: {p_rigged:.0%} win chance")
    ax.set_xticks(k)
    ax.set_xlabel(f"50/50s won out of {n}", color=TEXT, fontsize=24)
    ax.set_ylabel("% of players", color=TEXT, fontsize=24)
    ax.legend(facecolor=PANEL, edgecolor=GRID, labelcolor=TEXT, fontsize=22, loc="lower center", bbox_to_anchor=(0.5, 1.0), ncol=2, frameon=False)
    _save(fig, path)


def cost_of_n(results: dict, n: int, path: Path):
    fig, ax = plt.subplots(figsize=(16, 7.6))
    _style(ax, fig)
    names = [r["name"] for r in results.values()]
    p50 = [r["usd_p50"] for r in results.values()]
    p90 = [r["usd_p90"] for r in results.values()]
    y = np.arange(len(names))
    ax.barh(y + 0.2, p50, height=0.38, color=CYAN, label="typical player")
    ax.barh(y - 0.2, p90, height=0.38, color=RED, label="unlucky 1 in 10")
    for i, (a, b) in enumerate(zip(p50, p90)):
        ax.text(a + 30, i + 0.2, f"${a:,.0f}", color=TEXT, va="center", fontsize=22)
        ax.text(b + 30, i - 0.2, f"${b:,.0f}", color=TEXT, va="center", fontsize=22)
    ax.set_yticks(y, names, color=TEXT, fontsize=24)
    ax.set_xlabel(f"USD for {n} featured 5★ in a row, cash only", color=TEXT, fontsize=24)
    ax.set_xlim(0, max(p90) * 1.2)
    ax.legend(facecolor=PANEL, edgecolor=GRID, labelcolor=TEXT, fontsize=22, loc="lower center", bbox_to_anchor=(0.5, 1.0), ncol=2, frameon=False)
    _save(fig, path)


def pair_bars(results: dict, a_key: str, b_key: str, a_label: str, b_label: str, metric: str, xlabel: str, path: Path):
    """Two grouped bars per game, e.g. character only vs character + weapon (median or p90 USD)."""
    fig, ax = plt.subplots(figsize=(16, 7.6))
    _style(ax, fig)
    names = [r["name"] for r in results.values()]
    av = [r["builds"][a_key][metric] for r in results.values()]
    bv = [r["builds"][b_key][metric] for r in results.values()]
    y = np.arange(len(names))
    ax.barh(y + 0.2, av, height=0.38, color=CYAN, label=a_label)
    ax.barh(y - 0.2, bv, height=0.38, color=AMBER, label=b_label)
    for i, (x1, x2) in enumerate(zip(av, bv)):
        ax.text(x1 + 8, i + 0.2, f"${x1:,.0f}", color=TEXT, va="center", fontsize=22)
        ax.text(x2 + 8, i - 0.2, f"${x2:,.0f}", color=TEXT, va="center", fontsize=22)
    ax.set_yticks(y, names, color=TEXT, fontsize=24)
    ax.set_xlabel(xlabel, color=TEXT, fontsize=24)
    ax.set_xlim(0, max(bv) * 1.2)
    ax.legend(facecolor=PANEL, edgecolor=GRID, labelcolor=TEXT, fontsize=22, loc="lower center", bbox_to_anchor=(0.5, 1.0), ncol=2, frameon=False)
    _save(fig, path)


def ladder(results: dict, metric: str, ylabel: str, path: Path):
    """Cost of C0..C6 per game."""
    fig, ax = plt.subplots(figsize=(16, 7.6))
    _style(ax, fig)
    steps = [f"C{i}" for i in range(7)]
    ends = []
    for (g, r), c in zip(results.items(), PALETTE):
        vals = [r["ladder"][s][metric] for s in steps]
        ax.plot(range(7), vals, color=c, linewidth=4, marker="o", markersize=10, label=r["name"])
        ends.append([vals[-1], vals[-1], c])
    # keep end labels apart (Genshin and Star Rail finish within $15 of each other)
    ends.sort(key=lambda x: x[0])
    gap = (max(x[0] for x in ends) - min(r["ladder"]["C0"][metric] for r in results.values())) * 0.075
    for i in range(1, len(ends)):
        ends[i][1] = max(ends[i][1], ends[i - 1][1] + gap)
    for v, ypos, c in ends:
        ax.text(6.08, ypos, f"${v:,.0f}", color=c, va="center", fontsize=21)
    ax.set_xticks(range(7), ["0 copies extra\n(C0 / E0)" if i == 0 else f"C{i} / E{i}" for i in range(7)], color=TEXT, fontsize=19)
    ax.set_ylabel(ylabel, color=TEXT, fontsize=24)
    ax.set_xlim(-0.2, 6.9)
    ax.legend(facecolor=PANEL, edgecolor=GRID, labelcolor=TEXT, fontsize=20, loc="upper left", frameon=False)
    _save(fig, path)