#!/usr/bin/env python3
"""正式实验。三个互不依赖的部分，可以分进程并行：
  main     —— 3 种划分（1° 分块 / 2.5° 分块 / 像元随机）× 3 种子 × 仅重力 / 仅磁 / 早期融合（+ 后期融合）
  nulls    —— 标注置换零模型（融合）+ 环形平移零模型（三个模型），每种子 3 次
  rulers   —— 正 / 负尺子对照（合成标注），1° 分块，三个模型，3 种子
用法: python3 01_experiment.py <grid.npz> <part> <out_dir>
每个拟合都打一行日志；超参数全部固定，不调参。
"""
import sys, json, time
from pathlib import Path
import numpy as np
from sklearn.ensemble import HistGradientBoostingClassifier

GRID, PART, OUT = Path(sys.argv[1]), sys.argv[2], Path(sys.argv[3])
OUT.mkdir(parents=True, exist_ok=True)
d = np.load(GRID, allow_pickle=True)
X, y, iy, ix = d["X"].astype("float64"), d["y"].astype(int), d["iy"].astype(int), d["ix"].astype(int)
LAB = d["label_grid"].astype(int)
NY, NX = LAB.shape
SEEDS = [0, 1, 2]
COLS = {"gravity_only": list(range(0, 4)), "magnetic_only": list(range(4, 8)), "early_fusion": list(range(0, 8))}
T0 = time.time()

def log(msg):
    print(f"[{time.time()-T0:7.1f}s] {msg}", flush=True)

def folds(kind, seed, iy, ix):
    rng = np.random.default_rng(seed)
    if kind == "random_pixel":
        f = np.empty(iy.size, dtype=int); f[rng.permutation(iy.size)] = np.arange(iy.size) % 5
        return f, (iy // 30) * 1000 + (ix // 30)
    size = {"block_1deg": 30, "block_2p5deg": 75}[kind]
    blk = (iy // size) * 1000 + (ix // size)
    ub = np.unique(blk)
    fb = np.empty(ub.size, dtype=int); fb[rng.permutation(ub.size)] = np.arange(ub.size) % 5
    return fb[np.searchsorted(ub, blk)], (iy // 30) * 1000 + (ix // 30)

def bacc(yt, yp, k=4):
    rec = [np.mean(yp[yt == c] == c) for c in range(k) if np.any(yt == c)]
    return float(np.mean(rec))

def cv(Xm, yv, f, seed, proba=False):
    pred = np.empty(yv.size, dtype=int); P = np.zeros((yv.size, 4)) if proba else None
    per_fold = []
    for k in range(5):
        te = f == k; tr = ~te
        m = HistGradientBoostingClassifier(max_iter=200, class_weight="balanced", random_state=seed)
        m.fit(Xm[tr], yv[tr])
        if proba:
            pp = np.zeros((te.sum(), 4)); pp[:, m.classes_] = m.predict_proba(Xm[te]); P[te] = pp
            pred[te] = pp.argmax(1)
        else:
            pred[te] = m.predict(Xm[te])
        per_fold.append(bacc(yv[te], pred[te]))
    return pred, P, per_fold

res = {"part": PART, "seeds": SEEDS}
if PART == "main":
    save = {"y": y, "iy": iy, "ix": ix}
    for kind in ("block_1deg", "block_2p5deg", "random_pixel"):
        for s in SEEDS:
            f, blk1 = folds(kind, s, iy, ix)
            save[f"fold_{kind}_s{s}"] = f
            probs = {}
            for arm, cols in COLS.items():
                pred, P, pf = cv(X[:, cols], y, f, s, proba=True)
                save[f"pred_{kind}_{arm}_s{s}"] = pred; probs[arm] = P
                res.setdefault(kind, {}).setdefault(arm, {})[str(s)] = {"ba": bacc(y, pred), "per_fold_ba": pf}
                log(f"{kind} seed{s} {arm} BA={bacc(y, pred):.4f} folds={np.round(pf,3).tolist()}")
            late = ((probs["gravity_only"] + probs["magnetic_only"]) / 2).argmax(1)
            save[f"pred_{kind}_late_fusion_s{s}"] = late
            res[kind].setdefault("late_fusion", {})[str(s)] = {"ba": bacc(y, late)}
            log(f"{kind} seed{s} late_fusion BA={bacc(y, late):.4f}")
    np.savez_compressed(OUT / "oof_main.npz", **save)

elif PART == "nulls":
    Xg = np.full((NY, NX, 8), np.nan); Xg[iy, ix] = X
    for s in SEEDS:
        rng = np.random.default_rng(1000 + s)
        f, _ = folds("block_1deg", s, iy, ix)
        for r in range(3):
            yp = y[rng.permutation(y.size)]
            pred, _, _ = cv(X, yp, f, s)
            res.setdefault("label_permutation", {}).setdefault("early_fusion", {}).setdefault(str(s), []).append(bacc(yp, pred))
            log(f"perm seed{s} rep{r} early_fusion BA={bacc(yp, pred):.4f}")
        for r in range(3):
            dy = int(rng.integers(90, NY - 90)); dx = int(rng.integers(90, NX - 90))
            L = np.roll(np.roll(LAB, dy, 0), dx, 1)
            ok = (L >= 0) & np.all(np.isfinite(Xg), -1)
            jy, jx = np.nonzero(ok); Xs = Xg[jy, jx]; ys = L[jy, jx]
            fs, _ = folds("block_1deg", s, jy, jx)
            for arm, cols in COLS.items():
                pred, _, _ = cv(Xs[:, cols], ys, fs, s)
                res.setdefault("toroidal_shift", {}).setdefault(arm, {}).setdefault(str(s), []).append(
                    {"ba": bacc(ys, pred), "shift_cells": [dy, dx], "n": int(ys.size)})
                log(f"shift seed{s} rep{r} ({dy},{dx}) {arm} BA={bacc(ys, pred):.4f}")

elif PART == "rulers":
    save = {"y": y, "iy": iy, "ix": ix}
    g, m = X[:, 3], X[:, 7]                       # grav_resid15, mag_resid15
    for s in SEEDS:
        rng = np.random.default_rng(2000 + s)
        f, _ = folds("block_1deg", s, iy, ix)
        save[f"fold_block_1deg_s{s}"] = f
        flip = rng.random(y.size) < 0.25
        rand = rng.integers(0, 4, y.size)
        pos = 2 * (g > np.median(g)) + (m > np.median(m))
        neg = np.digitize(g, np.quantile(g, [0.25, 0.5, 0.75]))
        for name, yy in (("positive", pos), ("negative", neg)):
            yy = np.where(flip, rand, yy).astype(int)
            save[f"y_{name}_s{s}"] = yy
            for arm, cols in COLS.items():
                pred, _, pf = cv(X[:, cols], yy, f, s)
                save[f"pred_{name}_{arm}_s{s}"] = pred
                res.setdefault(name, {}).setdefault(arm, {})[str(s)] = {"ba": bacc(yy, pred), "per_fold_ba": pf}
                log(f"ruler {name} seed{s} {arm} BA={bacc(yy, pred):.4f}")
    np.savez_compressed(OUT / "oof_rulers.npz", **save)

res["wall_seconds"] = round(time.time() - T0, 1)
(OUT / f"raw_{PART}.json").write_text(json.dumps(res, indent=1))
log(f"done {PART}")