#!/usr/bin/env python3
"""补充分析（跑完之后加的，不属于预注册判据，只作说明）：
在像元随机划分与 2.5° 分块下，同样按 1° 块重抽 / 按块互换，给「早期融合 − 仅重力」的 ΔBA 配上区间与 p，
用来量化：如果用随机划分做评估，会得出一个多大的、实际上不存在的融合增益。
复用 02_stats.py 里完全相同的 paired() 实现。
用法: python3 02b_supplement.py <results_dir>
"""
import sys, json
from pathlib import Path
import numpy as np

R = Path(sys.argv[1])
main = np.load(R / "oof_main.npz")
y, iy, ix = main["y"], main["iy"], main["ix"]
blk1 = (iy // 30) * 1000 + (ix // 30)
SEEDS = [0, 1, 2]; B = 1000
P = {c: main[c] for c in main.files if c.startswith("pred_")}

def conf_by_block(yt, yp, blk):
    ub, inv = np.unique(blk, return_inverse=True)
    C = np.zeros((ub.size, 4, 4)); np.add.at(C, (inv, yt, yp), 1); return C

def ba_from_conf(C):
    tot = C.sum(-1); diag = np.diagonal(C, axis1=-2, axis2=-1)
    with np.errstate(invalid="ignore", divide="ignore"):
        rec = np.where(tot > 0, diag / np.where(tot > 0, tot, 1), np.nan)
    return np.nanmean(rec, -1)

def paired(yt, pa, pb, blk, seed):
    CA = conf_by_block(yt, pa, blk); CB = conf_by_block(yt, pb, blk)
    d_obs = float(ba_from_conf(CA.sum(0)) - ba_from_conf(CB.sum(0)))
    rng = np.random.default_rng(5000 + seed)
    W = rng.multinomial(CA.shape[0], np.full(CA.shape[0], 1 / CA.shape[0]), size=B)
    boot = ba_from_conf(np.einsum("bk,kij->bij", W, CA)) - ba_from_conf(np.einsum("bk,kij->bij", W, CB))
    S = rng.random((B, CA.shape[0])) < 0.5
    SA = np.einsum("bk,kij->bij", S, CB) + np.einsum("bk,kij->bij", ~S, CA)
    SB = np.einsum("bk,kij->bij", S, CA) + np.einsum("bk,kij->bij", ~S, CB)
    sw = ba_from_conf(SA) - ba_from_conf(SB)
    p = float((1 + np.sum(np.abs(sw) >= abs(d_obs) - 1e-12)) / (B + 1))
    lo, hi = np.percentile(boot, [2.5, 97.5])
    return {"delta_ba": d_obs, "ci95": [float(lo), float(hi)], "p_block_swap": p}

out = {"_note": "跑完之后加的补充分析，不参与 P1–P4 判定。", "early_fusion_minus_gravity_only": {}}
for kind in ("random_pixel", "block_1deg", "block_2p5deg"):
    per = {str(s): paired(y, P[f"pred_{kind}_early_fusion_s{s}"], P[f"pred_{kind}_gravity_only_s{s}"], blk1, s) for s in SEEDS}
    out["early_fusion_minus_gravity_only"][kind] = {"by_seed": per,
        "delta_ba_mean": float(np.mean([v["delta_ba"] for v in per.values()]))}
# 每类召回率（1° 分块，三种子平均），看融合 / 重力各自在哪一类上有本事
rec = {}
for kind in ("block_1deg", "random_pixel"):
    for arm in ("gravity_only", "magnetic_only", "early_fusion"):
        rs = []
        for s in SEEDS:
            C = conf_by_block(y, P[f"pred_{kind}_{arm}_s{s}"], blk1).sum(0)
            rs.append((np.diag(C) / C.sum(1)).tolist())
        rec.setdefault(kind, {})[arm] = np.mean(rs, 0).tolist()
out["per_class_recall_mean"] = {"classes": ["cenozoic_cover", "mesozoic_paleozoic_strata", "precambrian", "igneous"], **rec}
(R / "supplement.json").write_text(json.dumps(out, ensure_ascii=False, indent=1))
print(json.dumps(out, ensure_ascii=False, indent=1))