#!/usr/bin/env python3
"""Publication-grade figures for the AIP-AD experiments."""
from __future__ import annotations

import json
import os
import sys

import matplotlib
matplotlib.use("Agg")
import matplotlib.font_manager as fm
for f in ("/usr/share/fonts/truetype/chinese/NotoSansSC-Regular.ttf",
          "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"):
    if os.path.exists(f):
        fm.fontManager.addfont(f)
import matplotlib.pyplot as plt
import numpy as np

plt.rcParams["font.sans-serif"] = ["DejaVu Sans", "Noto Sans SC"]
plt.rcParams["axes.unicode_minus"] = False
plt.rcParams["figure.dpi"] = 150
plt.rcParams["savefig.bbox"] = "tight"

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
RES = os.path.join(ROOT, "results")
FIG = os.path.join(RES, "figures")
os.makedirs(FIG, exist_ok=True)
sys.path.insert(0, ROOT)


def style_ax(ax):
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.spines["left"].set_color("#cbd5e1")
    ax.spines["bottom"].set_color("#cbd5e1")
    ax.tick_params(colors="#475569", labelsize=8)
    ax.grid(True, linestyle="--", linewidth=0.5, alpha=0.2)
    ax.set_axisbelow(True)

C_MAX = "#94a3b8"
C_RAW = "#0891b2"
C_ACL = "#7c3aed"
C_RL = "#dc2626"
C_AIP = "#f59e0b"

with open(os.path.join(RES, "results.json")) as f:
    R = json.load(f)

cats = [c for c in R["fixed_strategies"]["max"] if c != "macro"]
table = R["fixed_strategies"]
acl_table = R["acl_strategies"]
arb = R["grpo_arbiter"]


def grouped_bars(ax, src, names, colors, title, arbiter_line=None):
    x = np.arange(len(cats) + 1)   # last slot = MACRO
    w = 0.8 / len(names)
    for i, n in enumerate(names):
        v = [src[n][c] for c in cats] + [src[n]["macro"]]
        ax.bar(x + i * w, v, w, label=n, color=colors[i], edgecolor="white", linewidth=0.4)
    if arbiter_line is not None:
        ax.axhline(arbiter_line, color=C_RL, ls="--", lw=1.6)
        ax.text(0.15, arbiter_line + 0.004, f"GRPO arbiter {arbiter_line:.3f}",
                color=C_RL, fontsize=8, weight="bold")
    ax.set_xticks(x + w * (len(names) - 1) / 2)
    ax.set_xticklabels(cats + ["MACRO"], rotation=20, ha="right", fontsize=8)
    ax.set_ylim(0.5, 1.0)
    ax.set_ylabel("Sample-level AUROC")
    ax.set_title(title, fontsize=10, weight="bold")
    style_ax(ax)
    ax.legend(fontsize=8, ncol=len(names), loc="lower left", bbox_to_anchor=(0, 1.01), frameon=False)


fig, axes = plt.subplots(1, 2, figsize=(13, 4.4), constrained_layout=True)
grouped_bars(axes[0], table, ["max", "mean", "median"], [C_MAX, C_RAW, C_RAW],
             "RAW camera agents — paper protocol (max) loses to robust/mean pooling")
grouped_bars(axes[1], acl_table, ["max", "median", "trimmed"], [C_MAX, C_ACL, C_ACL],
             "ACL-calibrated agents — robust poolers recover; arbiter tops all",
             arbiter_line=arb["per_category"]["macro"])
fig.savefig(os.path.join(FIG, "fig1_benchmark.png"))
plt.close(fig)

# --- fig2: radar of strategies (ACL) ---
names = ["max", "mean", "median", "top2", "trimmed", "aip_d05", "aip_c03", "aip_mt01", "aip_mt025", "aip_wt05"]
ang = np.linspace(0, 2 * np.pi, len(cats), endpoint=False).tolist()
ang += ang[:1]
fig, ax = plt.subplots(figsize=(5.8, 5.4), subplot_kw=dict(polar=True), constrained_layout=True)
pal = ["#64748b", "#0891b2", "#7c3aed", "#059669", "#65a30d", "#f59e0b", "#ea580c", "#db2777", "#be123c", "#dc2626"]
for n, c in zip(names, pal):
    v = [acl_table[n][c] for c in cats]
    v += v[:1]
    ax.plot(ang, v, label=n, color=c, linewidth=1.2)
    ax.fill(ang, v, color=c, alpha=0.04)
av = [arb["per_category"][c] for c in cats]
av += av[:1]
ax.plot(ang, av, label="GRPO arbiter", color="#dc2626", linewidth=2.4, linestyle="--")
ax.set_xticks(ang[:-1])
ax.set_xticklabels(cats, fontsize=8)
ax.set_ylim(0.5, 1.0)
ax.set_title("Fusion strategies on ACL-calibrated agents (S-AUROC)", fontsize=10, weight="bold", pad=18)
ax.legend(fontsize=6.5, loc="upper right", bbox_to_anchor=(1.34, 1.12))
fig.savefig(os.path.join(FIG, "fig2_radar.png"))
plt.close(fig)

# --- fig3: continual learning + online hedge ---
fig, axes = plt.subplots(1, 2, figsize=(11.5, 4.2), constrained_layout=True)
for ci, run in enumerate(R["continual"]):
    xs = np.arange(1, len(run["curve"]) + 1)
    c = ["#2563eb", "#dc2626"][ci]
    axes[0].plot(xs, [q["macro_sauroc_all6"] for q in run["curve"]], "o-", color=c,
                 label=f"order {ci+1} (all 6 cats)")
    axes[0].plot(xs, [q["macro_sauroc_seen"] for q in run["curve"]], "o--", color=c, alpha=0.45,
                 label=f"order {ci+1} (seen cats)")
for src, lab, col in ((acl_table["max"], "static max-pool (paper protocol)", C_MAX),
                      (acl_table["median"], "static median (ACL)", C_ACL)):
    axes[0].axhline(np.mean([src[c2] for c2 in cats]), color=col, ls=":", lw=1.5)
    axes[0].text(1.05, np.mean([src[c2] for c2 in cats]) + 0.003, lab, color=col, fontsize=7)
axes[0].set_xticks(range(1, len(cats) + 1))
axes[0].set_xlabel("stream step (category)")
axes[0].set_ylabel("macro S-AUROC")
axes[0].set_title("Continual learning of the GRPO arbiter (solid: all cats, dashed: seen)", fontsize=9, weight="bold")
axes[0].legend(fontsize=7, loc="lower left", bbox_to_anchor=(0, 1.02), ncol=2, frameon=False)
style_ax(axes[0])

hc = R["online_hedge"]
xs = np.arange(1, len(hc) + 1)
axes[1].plot(xs, [h["macro_sauroc_all"] for h in hc], "o-", color=C_AIP, label="online Hedge (operator feedback)")
axes[1].axhline(np.mean([table["max"][c] for c in cats]), color=C_MAX, ls=":", lw=1.5, label="static max-pool (raw)")
axes[1].set_xticks(xs)
axes[1].set_xticklabels([h["after"] for h in hc], rotation=20, ha="right", fontsize=8)
axes[1].set_ylabel("macro S-AUROC (cumulative)")
axes[1].set_title("Online learning: multiplicative-weights fusion", fontsize=9, weight="bold")
axes[1].legend(fontsize=8, loc="lower left", bbox_to_anchor=(0, 1.02), frameon=False)
style_ax(axes[1])
fig.savefig(os.path.join(FIG, "fig3_continual.png"))
plt.close(fig)

# --- fig4: policy behaviour ---
fig, axes = plt.subplots(1, 2, figsize=(11.5, 4.2), constrained_layout=True)
hist = arb["train_history"]
ep = [h["epoch"] for h in hist]
axes[0].plot(ep, [h["policy_entropy"] for h in hist], "-", color=C_RL)
axes[0].set_xlabel("GRPO epoch")
axes[0].set_ylabel("policy entropy (nats)")
axes[0].set_title("Policy specialization during GRPO training", fontsize=9, weight="bold")
style_ax(axes[0])
P = np.array(hist[-1]["mean_probs"])
snames = ["max", "mean", "median", "top2", "trimmed", "aip_d05", "aip_c03", "aip_soft2",
          "aip_mt01", "aip_mt025", "aip_mt05", "aip_mt075", "aip_wt05", "aip_wt075"]
order = np.argsort(P)
axes[1].barh(np.array(snames)[order], P[order], color=C_RL, alpha=0.85)
axes[1].set_xlabel("mean policy weight")
axes[1].set_title("Learned policy over pooling strategies (after ACL)", fontsize=9, weight="bold")
style_ax(axes[1])
axes[1].grid(axis="x", linestyle="--", linewidth=0.5, alpha=0.2)
fig.savefig(os.path.join(FIG, "fig4_policy.png"))
plt.close(fig)

# --- fig5: why dissent is information (swarm statistics, ACL scale) ---
gaps_ok, gaps_ng, top1_ok, top1_ng = [], [], [], []
for cat in cats:
    ev = np.load(os.path.join(RES, "scores", f"{cat}.npz.eval.npz"))
    vs = ev["view_scores"].astype(np.float64)
    y = ev["view_labels"].max(1)
    srt = np.sort(vs, axis=1)
    for i in np.where(y == 0)[0]:
        gaps_ok.append(srt[i, -1] - srt[i, -2])
        top1_ok.append(srt[i, -1])
    for i in np.where(y == 1)[0]:
        gaps_ng.append(srt[i, -1] - srt[i, -2])
        top1_ng.append(srt[i, -1])
fig, axes = plt.subplots(1, 2, figsize=(11.5, 4.2), constrained_layout=True)
axes[0].hist(gaps_ok, bins=48, alpha=0.62, label=f"normal (n={len(gaps_ok)})", color=C_MAX, density=True)
axes[0].hist(gaps_ng, bins=48, alpha=0.62, label=f"anomalous (n={len(gaps_ng)})", color=C_RL, density=True)
axes[0].set_xlabel("dissent: top-1 minus top-2 per-view score")
axes[0].set_ylabel("density")
axes[0].set_title("The dissenting 'liar' view carries the anomaly signal", fontsize=9, weight="bold")
axes[0].legend(fontsize=8, loc="lower left", bbox_to_anchor=(0, 1.02), frameon=False)
style_ax(axes[0])
axes[1].hist(top1_ok, bins=48, alpha=0.62, label="normal", color=C_MAX, density=True)
axes[1].hist(top1_ng, bins=48, alpha=0.62, label="anomalous", color=C_RL, density=True)
axes[1].set_xlabel("top-1 per-view score (max report)")
axes[1].set_title("Max report alone is weakly separated", fontsize=9, weight="bold")
axes[1].legend(fontsize=8, loc="lower left", bbox_to_anchor=(0, 1.02), frameon=False)
style_ax(axes[1])
fig.savefig(os.path.join(FIG, "fig5_dissent.png"))
plt.close(fig)

print("figures written to", FIG)
for f in sorted(os.listdir(FIG)):
    print("  ", f, os.path.getsize(os.path.join(FIG, f)))