#!/usr/bin/env python3
"""出图（只读结果文件，不做新计算）。
用法: python3 04_figs.py <results_dir> <grid.npz> <figs_dir>
"""
import sys, json
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

R, GRID, F = Path(sys.argv[1]), Path(sys.argv[2]), Path(sys.argv[3])
F.mkdir(parents=True, exist_ok=True)
e = json.loads((R / "experiment.json").read_text())
sup = json.loads((R / "supplement.json").read_text())
plt.rcParams.update({"font.size": 10, "axes.spines.top": False, "axes.spines.right": False})

# ---- 图 1：三种划分下各臂 BA + 零模型带 ----
arms = [("gravity_only", "Gravity only"), ("magnetic_only", "Magnetic only"),
        ("early_fusion", "Early fusion"), ("late_fusion", "Late fusion")]
kinds = [("random_pixel", "Random pixel split\n(leaky)"), ("block_1deg", "1° spatial blocks\n(main)"),
         ("block_2p5deg", "2.5° spatial blocks\n(conservative)")]
cols = ["#4C72B0", "#DD8452", "#55A868", "#8172B3"]
fig, ax = plt.subplots(figsize=(9, 4.6))
w = 0.2
for j, (k, kl) in enumerate(kinds):
    for i, (a, al) in enumerate(arms):
        v = e["arms"][k][a]["ba_by_seed"]
        x = j + (i - 1.5) * w
        ax.bar(x, np.mean(v), w * 0.9, color=cols[i], label=al if j == 0 else None)
        ax.scatter([x] * 3, v, s=8, color="k", zorder=3)
n = e["nulls"]
for key, c, lab in (("toroidal_shift_early_fusion", "#55A868", "Toroidal-shift null, fusion (max of 9)"),
                    ("toroidal_shift_gravity_only", "#4C72B0", "Toroidal-shift null, gravity (max of 9)"),
                    ("toroidal_shift_magnetic_only", "#DD8452", "Toroidal-shift null, magnetic (max of 9)")):
    ax.axhline(n[key]["max"], color=c, ls=":", lw=1.2, label=lab)
ax.axhline(0.25, color="grey", ls="--", lw=1, label="Majority class / chance = 0.25")
ax.set_xticks(range(3)); ax.set_xticklabels([kl for _, kl in kinds])
ax.set_ylabel("Balanced accuracy (4 classes)"); ax.set_ylim(0.2, 0.6)
ax.set_title("Same model, same data: the split decides whether fusion looks useful")
ax.legend(fontsize=7.5, ncol=2, loc="upper right", frameon=False)
fig.tight_layout(); fig.savefig(F / "fig1_ba_by_split.png", dpi=150); plt.close(fig)

# ---- 图 2：ΔBA（早期融合 − 仅重力）每种子 95% 区间，三种划分 + 两把尺子 ----
rows = []
for k, lab in (("random_pixel", "Random pixel split (post-hoc)"), ("block_1deg", "1° blocks (P1, main)"),
               ("block_2p5deg", "2.5° blocks (sensitivity)")):
    for s, v in sup["early_fusion_minus_gravity_only"][k]["by_seed"].items():
        rows.append((f"{lab} · seed {s}", v["delta_ba"], v["ci95"]))
for nm, lab in (("positive", "Positive ruler (label needs both)"), ("negative", "Negative ruler (label = gravity only)")):
    for s, v in e["rulers"][nm]["by_seed"].items():
        rows.append((f"{lab} · seed {s}", v["delta_ba"], v["ci95"]))
fig, ax = plt.subplots(figsize=(10, 5.4))
yy = np.arange(len(rows))[::-1]
for yv, (lab, d, ci) in zip(yy, rows):
    c = "#C44E52" if "ruler" in lab else "#333333"
    ax.plot(ci, [yv, yv], color=c, lw=2); ax.scatter([d], [yv], color=c, s=18, zorder=3)
ax.axvline(0, color="grey", lw=1); ax.axvspan(-0.02, 0.02, color="#DDDDDD", alpha=0.6, label="±0.02 (P1 'no practical gain')")
ax.set_yticks(yy); ax.set_yticklabels([r[0] for r in rows], fontsize=7.5)
ax.set_xlabel("ΔBA = early fusion − gravity only (95% CI, 1° block bootstrap)")
ax.set_title("Fusion gain: ruler detects +0.36 when it exists; real data shows ~0")
ax.legend(fontsize=8, frameon=False, loc="lower right")
fig.subplots_adjust(left=0.36, right=0.97, top=0.93, bottom=0.1); fig.savefig(F / "fig2_delta_ba_ci.png", dpi=150); plt.close(fig)

# ---- 图 3：地图——标注与 1° 分块下融合模型的预测（种子 0） ----
d = np.load(GRID, allow_pickle=True)
m = np.load(R / "oof_main.npz")
LAB = d["label_grid"].astype(float); LAB[LAB < 0] = np.nan
from matplotlib.colors import ListedColormap
cmap = ListedColormap(["#E8D27A", "#6BAED6", "#9E6A9E", "#D6604D"])
P = np.full(LAB.shape, np.nan); P[m["iy"], m["ix"]] = m["pred_block_1deg_early_fusion_s0"]
G = d["grav_grid"].astype(float); M = d["mag_grid"].astype(float)
ext = [105, 120, 35, 50]
fig, axs = plt.subplots(1, 4, figsize=(15, 4.2))
for a_, img, t, cm in ((axs[0], G, "GGMplus gravity disturbance (mGal)", "RdBu_r"),
                       (axs[1], M, "EMAG2v3 UpCont magnetic anomaly (nT)", "RdBu_r"),
                       (axs[2], LAB, "Label: USGS geo3al → 4 classes", cmap),
                       (axs[3], P, "Early fusion prediction, 1° blocks, seed 0", cmap)):
    if cm is cmap:
        im = a_.imshow(img, origin="lower", extent=ext, cmap=cmap, vmin=-0.5, vmax=3.5, interpolation="nearest")
    else:
        lo, hi = np.nanpercentile(img, [2, 98]); r_ = max(abs(lo), abs(hi))
        im = a_.imshow(img, origin="lower", extent=ext, cmap=cm, vmin=-r_, vmax=r_)
        fig.colorbar(im, ax=a_, fraction=0.046, pad=0.02)
    a_.set_title(t, fontsize=9); a_.set_xlabel("Lon (°E)")
axs[0].set_ylabel("Lat (°N)")
from matplotlib.patches import Patch
fig.legend(handles=[Patch(color=cmap(i), label=l) for i, l in enumerate(
    ["Cenozoic cover", "Mesozoic–Paleozoic strata", "Precambrian", "Igneous"])],
    loc="lower center", ncol=4, frameon=False, fontsize=9)
fig.tight_layout(rect=[0, 0.07, 1, 1]); fig.savefig(F / "fig3_maps.png", dpi=130); plt.close(fig)
print("figs written:", sorted(p.name for p in F.glob("*.png")))