#!/usr/bin/env python3
"""出图 + 策略轨迹报告（明确要的产物）。"""
from __future__ import annotations

import json
from collections import Counter
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

plt.rcParams["font.sans-serif"] = ["Noto Sans CJK JP", "Noto Serif CJK SC",
                                   "WenQuanYi Zen Hei Sharp", "DejaVu Sans"]
plt.rcParams["axes.unicode_minus"] = False

PKG = Path(__file__).resolve().parents[2]
RAW = json.loads((PKG / "work/results/eval_raw.json").read_text(encoding="utf-8"))
ANA = json.loads((PKG / "work/results/analysis.json").read_text(encoding="utf-8"))
FIG = PKG / "figs"
FIG.mkdir(exist_ok=True)
SEEDS = [str(s) for s in RAW["eval_seeds"]]
REAL = RAW["real"]
MONTHS = None

LABEL = {"null_no_intervention": "零模型（不干预）", "random_budget_matched": "随机（预算匹配）",
         "degree_top_b": "度中心性优先", "extremity_top_b": "极端度贪心（主对照）",
         "extremity_double_budget": "极端度贪心·预算翻倍（正对照）",
         "offline_rl_fqi": "离线强化学习 FQI"}
COLOR = {"null_no_intervention": "#666666", "random_budget_matched": "#8c9eb2",
         "degree_top_b": "#c58b3a", "extremity_top_b": "#2f6f4f",
         "extremity_double_budget": "#9fc4ae", "offline_rl_fqi": "#b33a3a"}
ORDER = ["null_no_intervention", "random_budget_matched", "degree_top_b",
         "extremity_top_b", "extremity_double_budget", "offline_rl_fqi"]


def traj(arm):
    return np.array([REAL[arm][s]["trajectory"] for s in SEEDS])


def fig_trajectories():
    fig, ax = plt.subplots(figsize=(9, 5))
    n = traj("null_no_intervention").shape[1]
    xs = np.arange(1, n + 1)
    for a in ORDER:
        T = traj(a)
        ax.plot(xs, T.mean(0), label=LABEL[a], color=COLOR[a], lw=2)
        if a in ("null_no_intervention", "offline_rl_fqi", "extremity_top_b"):
            ax.fill_between(xs, T.mean(0) - T.std(0), T.mean(0) + T.std(0),
                            color=COLOR[a], alpha=0.12, lw=0)
    ax.set_xlabel(f"留出期月窗（{RAW['eval_windows'][0]} → {RAW['eval_windows'][1]}，共 {n} 个月）")
    ax.set_ylabel("极化方差 Var(x)")
    ax.set_title("留出期极化方差轨迹（12 个评估种子均值，阴影为 ±1 标准差）")
    ax.legend(fontsize=8, loc="upper right")
    ax.grid(alpha=0.25)
    fig.tight_layout()
    fig.savefig(FIG / "polarization_trajectories.png", dpi=150)
    plt.close(fig)


def fig_reduction():
    fig, ax = plt.subplots(figsize=(9, 4.6))
    nullv = np.array([REAL["null_no_intervention"][s]["P_end"] for s in SEEDS])
    arms = [a for a in ORDER if a != "null_no_intervention"]
    for i, a in enumerate(arms):
        pe = np.array([REAL[a][s]["P_end"] for s in SEEDS])
        rel = (nullv - pe) / nullv * 100
        ax.bar(i, np.median(rel), color=COLOR[a], alpha=0.75, width=0.6)
        ax.scatter(np.full(len(rel), i) + np.linspace(-0.16, 0.16, len(rel)), rel,
                   s=14, color="#222222", zorder=3)
    ax.set_xticks(range(len(arms)))
    ax.set_xticklabels([LABEL[a] for a in arms], fontsize=8)
    ax.set_ylabel("相对零模型的极化方差降幅 (%)")
    ax.set_title("留出期末极化方差降幅（柱=12 个种子的中位数，点=每个种子）")
    ax.axhline(0, color="#000000", lw=0.8)
    ax.grid(axis="y", alpha=0.25)
    fig.tight_layout()
    fig.savefig(FIG / "reduction_vs_null.png", dpi=150)
    plt.close(fig)


def fig_paired():
    h1 = ANA["tests"]["H1_primary_rl_vs_extremity_P_end"]
    d = np.array(h1["paired_diffs"])
    fig, ax = plt.subplots(figsize=(8, 4.2))
    ax.scatter(range(1, len(d) + 1), d, s=45,
               color=["#2f6f4f" if v > 0 else "#b33a3a" for v in d], zorder=3)
    ax.axhline(0, color="#000000", lw=1)
    ax.axhline(h1["median_paired_diff"], color="#b33a3a", ls="--", lw=1.2,
               label=f"配对差值中位数 = {h1['median_paired_diff']:+.5f}")
    ax.set_xticks(range(1, len(d) + 1))
    ax.set_xticklabels(RAW["eval_seeds"], fontsize=8)
    ax.set_xlabel("评估种子")
    ax.set_ylabel("极端度贪心 − 离线 RL\n（>0 表示 RL 的极化方差更低）")
    ax.set_title("主对照配对差值：离线 RL 对 极端度贪心\n"
                 f"{h1['count_statement']}；配对 Wilcoxon 双侧 p = {h1['wilcoxon_p_two_sided_exact']:.4f}；"
                 f"rank-biserial = {h1['rank_biserial']:+.3f}", fontsize=10)
    ax.legend(fontsize=8)
    ax.grid(alpha=0.25)
    fig.tight_layout()
    fig.savefig(FIG / "primary_paired_difference.png", dpi=150)
    plt.close(fig)


def fig_policy_trajectory():
    names = [str(x) for x in np.load(PKG / "work/data/temporal_graph.npz",
                                     allow_pickle=True)["names"]]
    n = len(REAL["offline_rl_fqi"][SEEDS[0]]["picks_per_window"])
    xs = np.arange(1, n + 1)
    rl = np.array([REAL["offline_rl_fqi"][s]["picks_per_window"] for s in SEEDS])
    ex = np.array([REAL["extremity_top_b"][s]["picks_per_window"] for s in SEEDS])
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4.4))
    ax1.plot(xs, rl.mean(0), color=COLOR["offline_rl_fqi"], lw=2, marker="o", ms=4,
             label="离线 RL 每窗出手节点数")
    ax1.plot(xs, ex.mean(0), color=COLOR["extremity_top_b"], lw=2, marker="s", ms=4,
             label="极端度贪心每窗出手节点数")
    ax1.set_xlabel("留出期月窗")
    ax1.set_ylabel("该窗被扰动的节点数")
    ax1.set_title("策略轨迹：预算在 16 个月里怎么花")
    ax1.legend(fontsize=8)
    ax1.grid(alpha=0.25)

    cnt = Counter()
    for s in SEEDS:
        for w in REAL["offline_rl_fqi"][s]["picks"]:
            cnt.update(w)
    top = cnt.most_common(15)[::-1]
    ax2.barh([names[i] for i, _ in top], [c for _, c in top], color=COLOR["offline_rl_fqi"],
             alpha=0.8)
    ax2.set_xlabel("被离线 RL 选中的累计次数（12 个种子 × 16 个月窗）")
    ax2.set_title("离线 RL 最常扰动的 15 个 subreddit")
    ax2.tick_params(axis="y", labelsize=7)
    ax2.grid(axis="x", alpha=0.25)
    fig.tight_layout()
    fig.savefig(FIG / "policy_trajectory.png", dpi=150)
    plt.close(fig)

    rl_cnt, ex_cnt = Counter(), Counter()
    for s in SEEDS:
        for w in REAL["offline_rl_fqi"][s]["picks"]:
            rl_cnt.update(w)
        for w in REAL["extremity_top_b"][s]["picks"]:
            ex_cnt.update(w)
    overlap = sum((rl_cnt & ex_cnt).values()) / max(1, sum(rl_cnt.values()))
    report = {
        "what": "策略轨迹报告：离线 RL 在留出期 16 个月窗上的出手时机、出手规模与被扰动节点",
        "eval_windows": RAW["eval_windows"],
        "rl_picks_per_window_mean": [float(v) for v in rl.mean(0)],
        "extremity_picks_per_window_mean": [float(v) for v in ex.mean(0)],
        "rl_total_spend_per_seed": {s: REAL["offline_rl_fqi"][s]["spend"] for s in SEEDS},
        "extremity_total_spend_per_seed": {s: REAL["extremity_top_b"][s]["spend"] for s in SEEDS},
        "rl_top20_nodes": [{"subreddit": names[i], "times_selected": c}
                           for i, c in rl_cnt.most_common(20)],
        "extremity_top20_nodes": [{"subreddit": names[i], "times_selected": c}
                                  for i, c in ex_cnt.most_common(20)],
        "node_overlap_fraction_rl_vs_extremity": round(overlap, 4),
        "n_distinct_nodes_touched_rl": len(rl_cnt),
        "n_distinct_nodes_touched_extremity": len(ex_cnt),
    }
    (PKG / "work/results/policy_trajectory.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=1), encoding="utf-8")


def main() -> int:
    fig_trajectories()
    fig_reduction()
    fig_paired()
    fig_policy_trajectory()
    print("figs:", sorted(p.name for p in FIG.glob("*.png")))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())