#!/usr/bin/env python3
"""从 OOF 预测算统计量：ΔBA、按 1° 空间块整块重抽的 95% 区间、按块互换的配对 p，
并按预注册 decision_rule 判定 P1–P4 与尺子对照。输出 work/results/experiment.json 与导出给独立复算的 CSV。
用法: python3 02_stats.py <results_dir>
"""
import sys, json, gzip, csv
from pathlib import Path
import numpy as np

R = Path(sys.argv[1])
main = np.load(R / "oof_main.npz"); rul = np.load(R / "oof_rulers.npz")
raw_main = json.loads((R / "raw_main.json").read_text())
raw_nulls = json.loads((R / "raw_nulls.json").read_text())
raw_rul = json.loads((R / "raw_rulers.json").read_text())
SEEDS = [0, 1, 2]; B = 1000

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):                      # C [..., 4, 4]
    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 macro_f1(C):
    tp = np.diagonal(C); fp = C.sum(0) - tp; fn = C.sum(1) - tp
    f1 = np.where(2 * tp + fp + fn > 0, 2 * tp / np.maximum(2 * tp + fp + fn, 1), 0)
    return float(f1.mean())

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)   # [B, nblk]
    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,
            "n_blocks": int(CA.shape[0]), "n_boot": B}

y = main["y"]; iy = main["iy"]; ix = main["ix"]
blk1 = (iy // 30) * 1000 + (ix // 30)
out = {"n_cells": int(y.size), "class_counts": np.bincount(y, minlength=4).tolist()}

# ---- 各臂 BA / macro-F1 ----
arms = ["gravity_only", "magnetic_only", "early_fusion", "late_fusion"]
table = {}
for kind in ("block_1deg", "block_2p5deg", "random_pixel"):
    for arm in arms:
        bas, f1s = [], []
        for s in SEEDS:
            C = conf_by_block(y, main[f"pred_{kind}_{arm}_s{s}"], blk1).sum(0)
            bas.append(float(ba_from_conf(C))); f1s.append(macro_f1(C))
        table.setdefault(kind, {})[arm] = {"ba_by_seed": bas, "ba_mean": float(np.mean(bas)),
                                           "macro_f1_by_seed": f1s, "macro_f1_mean": float(np.mean(f1s)),
                                           "per_fold_ba": {str(s): raw_main[kind][arm][str(s)].get("per_fold_ba") for s in SEEDS}}
    table[kind]["majority_class"] = {"ba_mean": 0.25, "note": "永远预测多数类，BA 按定义恒为 0.25"}
out["arms"] = table

# ---- P1：融合 vs 最强单源 ----
b1 = table["block_1deg"]
best = max(("gravity_only", "magnetic_only"), key=lambda a: b1[a]["ba_mean"])
out["best_single_source"] = best
p1 = {}
for fusion in ("early_fusion", "late_fusion"):
    for kind in ("block_1deg", "block_2p5deg"):
        per = {str(s): paired(y, main[f"pred_{kind}_{fusion}_s{s}"], main[f"pred_{kind}_{best}_s{s}"], blk1, s) for s in SEEDS}
        p1[f"{fusion}_vs_{best}_{kind}"] = {"by_seed": per, "delta_ba_mean": float(np.mean([per[k]["delta_ba"] for k in per]))}
out["fusion_vs_best_single"] = p1
m = p1[f"early_fusion_vs_{best}_block_1deg"]
dm = m["delta_ba_mean"]
gain = dm >= 0.02 and sum(v["ci95"][0] > 0 for v in m["by_seed"].values()) >= 2
harm = dm <= -0.02 and sum(v["ci95"][1] < 0 for v in m["by_seed"].values()) >= 2
verdict = {"P1_main": {"delta_ba_mean": dm, "falsified": bool(gain or harm),
                       "direction_if_falsified": "fusion_gain" if gain else ("fusion_harm" if harm else None)}}

# ---- P2：随机划分泄漏 ----
infl = {a: float(table["random_pixel"][a]["ba_mean"] - table["block_1deg"][a]["ba_mean"]) for a in arms}
out["random_pixel_inflation_vs_block_1deg"] = infl
verdict["P2_spatial_leakage"] = {"inflation_early_fusion": infl["early_fusion"],
                                 "status": "held" if infl["early_fusion"] >= 0.10 else ("partial" if infl["early_fusion"] >= 0.05 else "falsified")}

# ---- P3 / P4：零模型 ----
nulls = {}
perm = raw_nulls["label_permutation"]["early_fusion"]
nulls["label_permutation_early_fusion"] = {"by_seed": perm, "all": [v for s in SEEDS for v in perm[str(s)]]}
for arm in ("gravity_only", "magnetic_only", "early_fusion"):
    sh = raw_nulls["toroidal_shift"][arm]
    vals = [r["ba"] for s in SEEDS for r in sh[str(s)]]
    obs = b1[arm]["ba_mean"]
    nulls[f"toroidal_shift_{arm}"] = {"by_seed": sh, "mean": float(np.mean(vals)), "sd": float(np.std(vals, ddof=1)),
                                     "max": float(np.max(vals)), "observed_ba_mean": obs,
                                     "gap_vs_null_mean": float(obs - np.mean(vals)),
                                     "z": float((obs - np.mean(vals)) / np.std(vals, ddof=1)),
                                     "p_empirical": float((1 + sum(v >= obs for v in vals)) / (len(vals) + 1))}
out["nulls"] = nulls
p3_ok = True; p3_detail = {}
for arm in ("gravity_only", "early_fusion"):
    for s in SEEDS:
        obs = b1[arm]["ba_by_seed"][s]
        nv = [r["ba"] for r in raw_nulls["toroidal_shift"][arm][str(s)]]
        if arm == "early_fusion":
            nv = nv + perm[str(s)]
        ok = obs > max(nv) and obs - np.mean(nv) >= 0.05
        p3_detail[f"{arm}_s{s}"] = {"observed": obs, "null_max": max(nv), "gap_vs_null_mean": float(obs - np.mean(nv)), "ok": bool(ok)}
        p3_ok &= ok
verdict["P3_signal_exists"] = {"falsified": not p3_ok, "detail": p3_detail}
g4 = nulls["toroidal_shift_magnetic_only"]["gap_vs_null_mean"]
verdict["P4_magnetic_weak"] = {"gap_vs_shift_null_mean": g4, "falsified": bool(g4 >= 0.05)}

# ---- 尺子 ----
ry = rul["y"]; rblk = (rul["iy"] // 30) * 1000 + (rul["ix"] // 30)
rulers = {}
for name in ("positive", "negative"):
    rbest = max(("gravity_only", "magnetic_only"), key=lambda a: np.mean([raw_rul[name][a][str(s)]["ba"] for s in SEEDS]))
    per = {str(s): paired(rul[f"y_{name}_s{s}"], rul[f"pred_{name}_early_fusion_s{s}"], rul[f"pred_{name}_{rbest}_s{s}"], rblk, s) for s in SEEDS}
    rulers[name] = {"best_single": rbest, "by_seed": per,
                    "ba": {a: [raw_rul[name][a][str(s)]["ba"] for s in SEEDS] for a in ("gravity_only", "magnetic_only", "early_fusion")}}
pos_ok = all(v["delta_ba"] >= 0.10 and v["ci95"][0] > 0 for v in rulers["positive"]["by_seed"].values())
neg_ok = sum((v["ci95"][0] <= 0 <= v["ci95"][1]) or abs(v["delta_ba"]) < 0.02 for v in rulers["negative"]["by_seed"].values()) >= 2
rulers["positive_passed"] = bool(pos_ok); rulers["negative_passed"] = bool(neg_ok)
out["rulers"] = rulers
verdict["ruler_kill_switch"] = {"positive_passed": bool(pos_ok), "negative_passed": bool(neg_ok),
                                "P1_void": not (pos_ok and neg_ok)}
out["verdict"] = verdict
out["wall_seconds"] = {"main": raw_main["wall_seconds"], "nulls": raw_nulls["wall_seconds"], "rulers": raw_rul["wall_seconds"]}
(R / "experiment.json").write_text(json.dumps(out, ensure_ascii=False, indent=1))

# ---- 导出给纯标准库独立复算：主划分与随机划分的逐格预测 ----
with gzip.open(R / "oof_for_recheck.csv.gz", "wt", newline="") as fh:
    w = csv.writer(fh)
    cols = [f"pred_{k}_{a}_s{s}" for k in ("block_1deg", "random_pixel") for a in arms for s in SEEDS]
    w.writerow(["y", "block1"] + cols)
    M = np.column_stack([y, blk1] + [main[c] for c in cols]).astype(np.int64)   # 一次性解压，避免逐元素重读 npz
    w.writerows(M.tolist())
print(json.dumps(verdict, ensure_ascii=False, indent=1))
print("best_single:", best, "P1 detail:", json.dumps(m, indent=1))