#!/usr/bin/env python3
import json, os
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

HERE = os.path.dirname(os.path.abspath(__file__))
RES = os.path.join(HERE, "..", "results")
FIGS = os.path.join(HERE, "..", "..", "figs")
os.makedirs(FIGS, exist_ok=True)
plt.rcParams.update({"figure.dpi": 150, "font.size": 9, "axes.grid": True,
                     "grid.alpha": 0.25, "axes.spines.top": False, "axes.spines.right": False})

fs = json.load(open(os.path.join(RES, "final_stats.json"), encoding="utf-8"))
tk = json.load(open(os.path.join(RES, "topk_precision.json"), encoding="utf-8"))
an = json.load(open(os.path.join(RES, "analysis.json"), encoding="utf-8"))
z = np.load(os.path.join(RES, "scores.npz"))

# --- Fig 1: 零模型分布 vs 实测 ---
fig, axes = plt.subplots(1, 2, figsize=(9, 3.4))
for ax, nulls, obs, deg, title, p in (
        (axes[0], z["null_main"], fs["H1_main"]["auroc_gout_rwr"], fs["H1_main"]["auroc_degree_only"],
         "Main: gout-seeded RWR -> CVD genes", fs["H1_main"]["empirical_p"]),
        (axes[1], z["null_temporal"], fs["H3_temporal"]["auroc_gout_rwr"], fs["H3_temporal"]["auroc_degree_only"],
         "Temporal hold-out: pre-2022 seeds -> new CVD genes", fs["H3_temporal"]["empirical_p"])):
    ax.hist(nulls, bins=60, color="#b9c6d6", edgecolor="none", label="degree-matched random seeds (n=3000)")
    ax.axvline(obs, color="#c0392b", lw=2, label=f"observed AUROC={obs:.4f}")
    ax.axvline(deg, color="#2c7d59", lw=1.6, ls="--", label=f"degree-only baseline={deg:.4f}")
    ax.axvline(0.5, color="#888", lw=1, ls=":", label="chance = 0.5")
    ax.set_title(f"{title}\nempirical p = {p:.2g}", fontsize=9)
    ax.set_xlabel("AUROC"); ax.set_ylabel("null draws")
    ax.legend(fontsize=6.5, loc="upper left")
fig.tight_layout(); fig.savefig(os.path.join(FIGS, "fig1_null_distributions.png")); plt.close(fig)

# --- Fig 2: precision@k ---
ks = [int(k) for k in tk["precision_at_k"]["gout_rwr"].keys()]
fig, ax = plt.subplots(figsize=(5.6, 3.6))
g = [tk["precision_at_k"]["gout_rwr"][str(k)] for k in ks]
ax.plot(ks, [x["value"] for x in g], "o-", color="#c0392b", label="gout-seeded RWR")
ax.plot(ks, [x["null_mean"] for x in g], "s--", color="#7f8c8d", label="null mean (degree-matched)")
ax.fill_between(ks, [x["null_mean"] for x in g], [x["null_p95"] for x in g],
                color="#bdc3c7", alpha=.45, label="null mean -> 95th pct")
ax.plot(ks, [tk["precision_at_k"]["degree_only"][str(k)]["value"] for k in ks], "^-",
        color="#2c7d59", label="degree-only baseline")
ax.axhline(tk["prevalence_in_eval_universe"], color="#2980b9", ls=":", label="base rate (4.74%)")
for k in (50,):
    v = tk["precision_at_k"]["gout_rwr"][str(k)]
    ax.annotate(f"k=50: {v['value']:.1%} vs null {v['null_mean']:.1%}\nempirical p={v['empirical_p']:.2f} (n.s.)",
                xy=(50, v["value"]), xytext=(85, 0.088), fontsize=7,
                arrowprops=dict(arrowstyle="->", lw=.8))
ax.set_xscale("log"); ax.set_xticks(ks); ax.set_xticklabels(ks)
ax.set_xlabel("top-k"); ax.set_ylabel("fraction that are known CVD genes")
ax.set_title("The top of the ranking is NOT better than the null model", fontsize=9, pad=10)
ax.legend(fontsize=6.5)
fig.tight_layout(); fig.savefig(os.path.join(FIGS, "fig2_precision_at_k.png")); plt.close(fig)

# --- Fig 3: 配对折 ---
fig, axes = plt.subplots(1, 2, figsize=(9, 3.4))
for ax, rows, title, blk in ((axes[0], an["paired_folds_h2"], "H2: gout seeds -> CVD genes", fs["H2_paired_holdout"]),
                             (axes[1], an["positive_control_h0"], "H0 positive control: gout -> held-out gout", fs["H0_positive_control"])):
    x = np.arange(len(rows))
    ax.plot(x, [r["auroc_gout_rwr"] for r in rows], "o", color="#c0392b", label="gout seeds")
    ax.plot(x, [r["auroc_matched_null"] for r in rows], "s", color="#7f8c8d", label="degree-matched null")
    for i, r in enumerate(rows):
        ax.plot([i, i], [r["auroc_matched_null"], r["auroc_gout_rwr"]], color="#ccc", lw=.9, zorder=0)
    ax.axhline(0.5, color="#888", ls=":", lw=1)
    ax.set_xlabel("fold (3 seeds x 5 folds)"); ax.set_ylabel("AUROC")
    ax.set_title(f"{title}\n{blk['wins']}/{blk['n_pairs']} folds, median gain {blk['median_difference']:+.4f}, "
                 f"Wilcoxon p={blk['wilcoxon_p_two_sided']:.2g}", fontsize=8)
    ax.legend(fontsize=7)
fig.tight_layout(); fig.savefig(os.path.join(FIGS, "fig3_paired_folds.png")); plt.close(fig)

# --- Fig 4: 效应量对比 ---
fig, ax = plt.subplots(figsize=(6.2, 3.2))
items = [("Positive control\n(gout -> held-out gout)", fs["H0_positive_control"]["mean_difference"]),
         ("Main\n(gout -> CVD)", fs["H1_main"]["delta_vs_null_mean"]),
         ("Temporal hold-out\n(pre-2022 -> new CVD)", fs["H3_temporal"]["delta_vs_null_mean"]),
         ("No text-mining channel\n(gout -> CVD)", fs["sensitivity_no_textmining"]["delta_vs_null_mean"])]
cols = ["#2c7d59", "#c0392b", "#d98c34", "#8e6bb5"]
ax.bar([i[0] for i in items], [i[1] for i in items], color=cols, width=.6)
for i, (lbl, v) in enumerate(items):
    ax.text(i, v + .002, f"+{v:.4f}", ha="center", fontsize=8)
ax.set_ylabel("AUROC gain over degree-matched null")
ax.set_title("Effect sizes: the cross-disease signal is ~1/3 of the same-disease signal", fontsize=9)
ax.tick_params(axis="x", labelsize=7)
fig.tight_layout(); fig.savefig(os.path.join(FIGS, "fig4_effect_sizes.png")); plt.close(fig)
print("figs written:", sorted(os.listdir(FIGS)))