#!/usr/bin/env python3
"""独立复算：只用 Python 标准库，不用 numpy / sklearn。
从逐格预测 CSV 重新数一遍混淆矩阵，重算每个臂的 BA、ΔBA 与随机划分抬高量，写 work/results/recheck.json。
用法: python3 03_recheck.py <results_dir>
"""
import sys, csv, gzip, json
from pathlib import Path
R = Path(sys.argv[1])
SEEDS = (0, 1, 2); ARMS = ("gravity_only", "magnetic_only", "early_fusion", "late_fusion")
counts = {}
with gzip.open(R / "oof_for_recheck.csv.gz", "rt", newline="") as fh:
    rd = csv.reader(fh); header = next(rd)
    idx = {h: i for i, h in enumerate(header)}
    pred_cols = header[2:]
    for c in pred_cols:
        counts[c] = [[0] * 4 for _ in range(4)]
    for row in rd:
        t = int(row[0])
        for c in pred_cols:
            counts[c][t][int(row[idx[c]])] += 1
def ba(M):
    rec = [M[k][k] / sum(M[k]) for k in range(4) if sum(M[k]) > 0]
    return sum(rec) / len(rec)
out = {"implementation": "pure stdlib", "arms": {}}
for kind in ("block_1deg", "random_pixel"):
    for a in ARMS:
        v = [ba(counts[f"pred_{kind}_{a}_s{s}"]) for s in SEEDS]
        out["arms"].setdefault(kind, {})[a] = {"ba_by_seed": v, "ba_mean": sum(v) / 3}
b1 = out["arms"]["block_1deg"]
best = "gravity_only" if b1["gravity_only"]["ba_mean"] >= b1["magnetic_only"]["ba_mean"] else "magnetic_only"
out["best_single_source"] = best
d = [b1["early_fusion"]["ba_by_seed"][i] - b1[best]["ba_by_seed"][i] for i in range(3)]
out["delta_ba_early_fusion_block_1deg"] = {"by_seed": d, "mean": sum(d) / 3}
dl = [b1["late_fusion"]["ba_by_seed"][i] - b1[best]["ba_by_seed"][i] for i in range(3)]
out["delta_ba_late_fusion_block_1deg"] = {"by_seed": dl, "mean": sum(dl) / 3}
out["random_pixel_inflation_early_fusion"] = out["arms"]["random_pixel"]["early_fusion"]["ba_mean"] - b1["early_fusion"]["ba_mean"]
(R / "recheck.json").write_text(json.dumps(out, indent=1))
print(json.dumps(out, indent=1))
