"""Sequence networks on the RAW yearly Landsat time series, frozen backbone +
last-layer adaptation.

    python tools/nnlab/ts_nn.py [--smoke] [--only <substring>]

Step 1  tensor cache  lab.CACHE/ts/{city}_T.npy  [N, 42, 6] float32 DN, median of
        QA-valid observations per year (src/features.py make_flat_view(agg="median"):
        a year is kept only if every band has a valid median and Blue <= BLUE_MAX),
        NaN where no kept observation; rows aligned with the processed feature table.
        NOTE: the raw data covers 1984-2025 (42 years), as the feature pipeline uses.
Step 2  inputs        6 SR bands + NDVI, NDBI, MNDWI + observed mask = 10 ch x 42 steps,
        linear interpolation in time (edges nearest, as to_year_matrices), indices from
        the interpolated SR bands; per-city standardisation per channel. Variant "ry"
        subtracts the city's per-year median over pixels (9 continuous channels) first.
Step 3  backbones     tcnn, gru, tcnn_dann; class-balanced CE, AdamW(1e-3, wd 1e-4),
        batch 256, constant LR, brightness/offset augmentation on the 6 band channels.
        CPU budget: per (backbone, variant) ONE run on the lab.madrid_split source part,
        checkpoints at EPOCH_GRID = (5, 10, 15), the checkpoint with the best
        pseudo-target macro F1 IS the backbone (no retrain on all Madrid); lam for "ft"
        chosen on the same model. 1 seed. Default combos: tcnn_st, tcnn_ry,
        tcnn_dann_st; --extended adds tcnn_dann_ry, gru_st, gru_ry (~2.5 h more).
Step 4  heads         zs (n=0), ft (lab.finetune_last_layer), joint (lab.joint_head);
        probabilities averaged over seeds.
Methods: ts_{backbone}_{variant}_{head}.
"""
from __future__ import annotations

import argparse
import copy
import json
import sys
import time
import warnings
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import lab  # noqa: E402
import numpy as np  # noqa: E402
import config as C  # noqa: E402  (src/config.py, put on sys.path by lab)

T0 = time.time()
YEARS = np.arange(1984, 2026)            # 42 years present in both raw files
NY = len(YEARS)
BANDS = list(C.BASE_BANDS)
TS_DIR = lab.CACHE / "ts"
NN_DIR = lab.CACHE / "ts_nn"
for _d in (TS_DIR, NN_DIR):
    _d.mkdir(parents=True, exist_ok=True)

HEADS = ("zs", "ft", "joint")
EPOCH_GRID = (5, 10, 15)            # spec: (10, 20, 30) - halved to fit the CPU budget
LAM_GRID = (0.1, 1.0, 10.0, 100.0)
SEEDS = (0,)                         # spec: 2 seeds - cut to 1 to fit the CPU budget
DEFAULT_COMBOS = (("tcnn", "st"), ("tcnn", "ry"), ("tcnn_dann", "st"))
EXTENDED_COMBOS = (("tcnn_dann", "ry"), ("gru", "st"), ("gru", "ry"))
BATCH = 256
LR, WD = 1e-3, 1e-4
SHOT_WEIGHT = 30.0


def _rss():
    try:
        import psutil
        mi = psutil.Process().memory_info()
        return f" rss={mi.rss / 2**30:.2f}G peak={getattr(mi, 'peak_wset', mi.rss) / 2**30:.2f}G"
    except Exception:
        return ""


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


# --------------------------------------------------------------------------- #
# Step 1: tensor cache
# --------------------------------------------------------------------------- #
RAW = {"madrid": lab.ROOT / "data" / "raw" / "madrid_train.parquet",
       "amsterdam": lab.ROOT / "data" / "raw" / "amsterdam_data.parquet"}


def build_tensor(city: str, P: np.ndarray) -> np.ndarray:
    """P: (N, 2) px_key, py_key of the feature table (row order)."""
    import pyarrow.parquet as pq
    path = TS_DIR / f"{city}_T.npy"
    if path.exists():
        T = np.load(path)
        assert T.shape == (len(P), NY, 6)
        return T
    t0 = time.time()
    cols = ["px_key", "py_key", "year"] + [f"{b}_{k}" for k in (1, 2, 3) for b in BANDS] \
        + [f"qa_valid_{k}" for k in (1, 2, 3)]          # no label columns are read
    key_f = P[:, 0].astype(np.int64) * 10_000_000 + P[:, 1].astype(np.int64)
    assert P[:, 1].min() >= 0 and P[:, 1].max() < 10_000_000
    order = np.argsort(key_f)
    ks = key_f[order]
    assert (np.diff(ks) > 0).all(), "duplicate pixel keys in feature table"
    T = np.full((len(P), NY, 6), np.nan, np.float32)
    seen = np.zeros((len(P), NY), bool)
    f = pq.ParquetFile(RAW[city])
    n_rows = n_kept = 0
    for batch in f.iter_batches(batch_size=200_000, columns=cols):
        tb = {c: batch.column(c) for c in cols}
        n = batch.num_rows
        n_rows += n
        oks = [np.asarray(tb[f"qa_valid_{k}"].fill_null(False).to_numpy(zero_copy_only=False), bool)
               for k in (1, 2, 3)]
        med = np.empty((n, 6), np.float32)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", RuntimeWarning)
            for j, b in enumerate(BANDS):
                st = []
                for k in (1, 2, 3):
                    v = np.asarray(tb[f"{b}_{k}"].to_numpy(zero_copy_only=False), np.float32).copy()
                    v[~oks[k - 1]] = np.nan
                    st.append(v)
                med[:, j] = np.nanmedian(np.stack(st), axis=0)
        good = ~np.isnan(med).any(1)
        good &= ~(med[:, 0] > C.BLUE_MAX)                 # make_flat_view: keep Blue <= BLUE_MAX
        px = np.asarray(tb["px_key"].to_numpy(zero_copy_only=False), np.int64)
        py = np.asarray(tb["py_key"].to_numpy(zero_copy_only=False), np.int64)
        yr = np.asarray(tb["year"].to_numpy(zero_copy_only=False), np.int64)
        key = px * 10_000_000 + py
        pos = np.searchsorted(ks, key).clip(0, len(ks) - 1)
        hit = (ks[pos] == key) & good
        assert ((yr[hit] >= YEARS[0]) & (yr[hit] <= YEARS[-1])).all()
        r, c = order[pos[hit]], yr[hit] - YEARS[0]
        assert not seen[r, c].any(), "duplicate pixel-year"
        seen[r, c] = True
        T[r, c] = med[hit]
        n_kept += int(hit.sum())
        del batch, tb, med, st
    found = seen.any(1)
    assert found.all(), f"{(~found).sum()} feature-table pixels have no observation"
    np.save(path, T)
    log(f"tensor {city}: {n_rows:,} raw rows, {n_kept:,} pixel-years kept, "
        f"shape {T.shape}, observed {seen.mean():.3f}, {time.time() - t0:.0f}s")
    return T


# --------------------------------------------------------------------------- #
# Step 2: inputs
# --------------------------------------------------------------------------- #
def interp_time(T: np.ndarray) -> np.ndarray:
    """Linear interpolation over the year axis, nearest at the edges (per pixel, band)."""
    N, Y, B = T.shape
    out = T.copy()
    x = np.arange(Y, dtype=np.float32)
    for b in range(B):
        M = out[:, :, b]
        obs = ~np.isnan(M)
        # index of previous / next observed year for every cell
        idx = np.where(obs, np.arange(Y)[None, :], -1)
        prev = np.maximum.accumulate(idx, axis=1)
        idx2 = np.where(obs, np.arange(Y)[None, :], Y)
        nxt = np.minimum.accumulate(idx2[:, ::-1], axis=1)[:, ::-1]
        rows = np.arange(N)[:, None]
        pv = np.where(prev >= 0, M[rows, prev.clip(0)], np.nan)
        nv = np.where(nxt < Y, M[rows, nxt.clip(None, Y - 1)], np.nan)
        w = np.where((prev >= 0) & (nxt < Y) & (nxt > prev),
                     (x[None, :] - prev) / np.maximum(nxt - prev, 1), 0.0).astype(np.float32)
        lin = pv + w * (nv - pv)
        fill = np.where(np.isnan(pv), nv, np.where(np.isnan(nv), pv, lin))
        out[:, :, b] = np.where(obs, M, fill)
    assert not np.isnan(out).any()
    return out


def make_inputs(T: np.ndarray, variant: str) -> np.ndarray:
    """[N, Y, 6] DN -> [N, 10, Y] float32, standardised with this city's statistics."""
    obs = (~np.isnan(T).any(2)).astype(np.float32)          # year observed
    S = interp_time(T) * np.float32(C.SR_SCALE) + np.float32(C.SR_OFFSET)
    B, G, R, N_, S1 = S[..., 0], S[..., 1], S[..., 2], S[..., 3], S[..., 4]
    eps = np.float32(1e-6)
    ndvi = (N_ - R) / (N_ + R + eps)
    ndbi = (S1 - N_) / (S1 + N_ + eps)
    mndwi = (G - S1) / (G + S1 + eps)
    X = np.concatenate([S, ndvi[..., None], ndbi[..., None], mndwi[..., None], obs[..., None]], 2)
    del S
    X = np.clip(X, -5, 5)        # index denominators near 0 on bad pixels
    if variant == "ry":
        X[..., :9] -= np.median(X[..., :9], axis=0, keepdims=True)
    mu = X.reshape(-1, 10).mean(0, dtype=np.float64)
    sd = X.reshape(-1, 10).std(0, dtype=np.float64)
    sd[sd < 1e-8] = 1
    X = ((X - mu.astype(np.float32)) / sd.astype(np.float32)).astype(np.float32)
    return np.ascontiguousarray(X.transpose(0, 2, 1))       # [N, C, Y]


# --------------------------------------------------------------------------- #
# Step 3: backbones
# --------------------------------------------------------------------------- #
def build_net(kind: str):
    import torch
    nn = torch.nn

    class GradReverse(torch.autograd.Function):
        @staticmethod
        def forward(ctx, x, lam):
            ctx.lam = lam
            return x.view_as(x)

        @staticmethod
        def backward(ctx, g):
            return -ctx.lam * g, None

    class TempCNN(nn.Module):
        def __init__(self, c_in=10, T=NY, dann=False):
            super().__init__()
            layers, c = [], c_in
            for _ in range(3):
                layers += [nn.Conv1d(c, 64, 5, padding=2), nn.BatchNorm1d(64), nn.ReLU(), nn.Dropout(0.2)]
                c = 64
            self.conv = nn.Sequential(*layers)
            self.emb = nn.Sequential(nn.Flatten(), nn.Linear(64 * T, 128), nn.ReLU(),
                                     nn.Linear(128, 64), nn.ReLU())
            self.head = nn.Linear(64, 4)
            self.dom = nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 1)) if dann else None

        def embed(self, x):
            return self.emb(self.conv(x))

        def forward(self, x, lam=0.0):
            z = self.embed(x)
            dom = self.dom(GradReverse.apply(z, lam)) if self.dom is not None else None
            return self.head(z), dom

    class GRUNet(nn.Module):
        def __init__(self, c_in=10):
            super().__init__()
            self.gru = nn.GRU(c_in, 64, num_layers=2, bidirectional=True, batch_first=True)
            self.emb = nn.Sequential(nn.Linear(128, 64), nn.ReLU())
            self.head = nn.Linear(64, 4)
            self.dom = None

        def embed(self, x):
            _, h = self.gru(x.transpose(1, 2))               # h: [4, B, 64]
            return self.emb(torch.cat([h[-2], h[-1]], 1))

        def forward(self, x, lam=0.0):
            return self.head(self.embed(x)), None

    if kind == "tcnn":
        return TempCNN()
    if kind == "tcnn_dann":
        return TempCNN(dann=True)
    if kind == "gru":
        return GRUNet()
    raise ValueError(kind)


def augment(torch, xb, gen):
    """Per-sample brightness scale U(0.9,1.1) and offset N(0,0.1) on the 6 band channels."""
    B = xb.shape[0]
    s = torch.empty(B, 1, 1).uniform_(0.9, 1.1, generator=gen)
    o = torch.randn(B, 1, 1, generator=gen) * 0.1
    xb = xb.clone()
    xb[:, :6] = xb[:, :6] * s + o
    return xb


def train(kind, Xs, ys, Xt, epochs, seed, eval_fn=None, eval_at=(), max_rows=None, sched_epochs=None):
    """Returns (net, {epoch: eval result}, {epoch: state_dict copy}, sec/epoch).
    sched_epochs: length of the DANN lambda schedule (default `epochs`); a retrain to a
    selected checkpoint passes the selection run's length so it reproduces that model."""
    sched = sched_epochs or epochs
    torch = lab.torch_setup(seed)
    gen = torch.Generator().manual_seed(seed)
    net = build_net(kind)
    opt = torch.optim.AdamW(net.parameters(), lr=LR, weight_decay=WD)
    counts = np.bincount(ys, minlength=5)[1:]
    w = torch.tensor(len(ys) / (4 * np.maximum(counts, 1)), dtype=torch.float32)
    ce, bce = torch.nn.CrossEntropyLoss(weight=w), torch.nn.BCEWithLogitsLoss()
    S, Y = torch.from_numpy(Xs), torch.from_numpy(ys - 1).long()
    Tt = torch.from_numpy(Xt) if Xt is not None else None
    dann = kind.endswith("_dann")
    n = len(S) if max_rows is None else min(len(S), max_rows)
    steps = int(np.ceil(n / BATCH))
    evals, states, t_ep = {}, {}, []
    for ep in range(1, epochs + 1):
        te = time.time()
        net.train()
        perm = torch.randperm(len(S), generator=gen)[:n]
        for k in range(steps):
            i = perm[k * BATCH:(k + 1) * BATCH]
            if len(i) < 2:
                continue
            xb = augment(torch, S[i], gen)
            if dann:
                p = ((ep - 1) * steps + k) / (sched * steps)
                lam = float(2 / (1 + np.exp(-10 * p)) - 1)
                logit, dom_s = net(xb, lam)
                j = torch.randint(0, len(Tt), (len(i),), generator=gen)
                _, dom_t = net(augment(torch, Tt[j], gen), lam)
                loss = ce(logit, Y[i]) + bce(dom_s.squeeze(1), torch.zeros(len(i)))                     + bce(dom_t.squeeze(1), torch.ones(len(i)))
            else:
                logit, _ = net(xb, 0.0)
                loss = ce(logit, Y[i])
            opt.zero_grad()
            loss.backward()
            opt.step()
        t_ep.append((time.time() - te) * len(S) / n)      # extrapolated to a full epoch
        if ep in eval_at:
            states[ep] = copy.deepcopy(net.state_dict())
            if eval_fn is not None:
                evals[ep] = eval_fn(net)
            log(f"  {kind} seed{seed} epoch {ep}/{epochs} {t_ep[-1]:.1f}s/full-epoch"
                + (f" eval={evals[ep]:.4f}" if ep in evals else ""))
    net.eval()
    return net, evals, states, float(np.mean(t_ep))


def embed(net, X, bs=2048):
    import torch
    net.eval()
    out = []
    with torch.no_grad():
        for k in range(0, len(X), bs):
            out.append(net.embed(torch.from_numpy(X[k:k + bs])).numpy())
    return np.concatenate(out).astype(np.float32)


def last_layer(net):
    return (net.head.weight.detach().numpy().astype(np.float64).copy(),
            net.head.bias.detach().numpy().astype(np.float64).copy())


def proba(Z, Wt, b):
    return lab.softmax(Z.astype(np.float64) @ Wt.T + b)


def mf1(y, P):
    return lab.W.macro_f1(y, lab.CLASSES[P.argmax(1)])

# --------------------------------------------------------------------------- #
# Backbone per (backbone, variant): one run on the Madrid spatial source split,
# checkpoints at the epoch grid, pick the epoch on the pseudo-target, keep that
# checkpoint (budget: no retrain on all Madrid), then choose lam on the same model.
# --------------------------------------------------------------------------- #
def simulate_lam(Zt, yt, W0, b0, reps):
    lam_score = {}
    for lam in LAM_GRID:
        sc = []
        for n in (5, 50):
            for rep in range(reps):
                rng = np.random.default_rng(7_000 + 1000 * n + rep)
                sh = np.concatenate([rng.choice(np.flatnonzero(yt == c), min(n, int((yt == c).sum())),
                                                replace=False) for c in lab.CLASSES])
                rest = np.setdiff1d(np.arange(len(yt)), sh)
                Wf, bf = lab.finetune_last_layer(W0, b0, Zt[sh], yt[sh], lam)
                sc.append(mf1(yt[rest], proba(Zt[rest], Wf, bf)))
        lam_score[lam] = float(np.mean(sc))
    return lam_score


def backbone(kind, var, d, Xm, Xa, cfg):
    """Returns (selection dict, [(Zm, Za, W0, b0) per seed])."""
    tag = f"{kind}_{var}{cfg['tag']}"
    sel_path = NN_DIR / f"select_{tag}.json"
    src, tgt = lab.madrid_split(d["pm"])
    ym = d["ym"]
    grid = cfg["epoch_grid"]
    embs = []
    sel = json.loads(sel_path.read_text()) if sel_path.exists() else None
    for seed in cfg["seeds"]:
        path = NN_DIR / f"{tag}_s{seed}.npz"
        if path.exists() and sel is not None:
            z = np.load(path)
            embs.append((z["Zm"], z["Za"], z["W0"], z["b0"]))
            continue
        t0 = time.time()
        if sel is None:
            def ev(net):
                Wt, b = last_layer(net)
                return mf1(ym[tgt], proba(embed(net, Xm[tgt]), Wt, b))
            net, evals, states, sec = train(kind, Xm[src], ym[src], Xa, max(grid), seed, eval_fn=ev,
                                            eval_at=grid, max_rows=cfg["max_rows"])
            E = max(grid, key=lambda e: (evals[e], -e))          # ties -> fewer epochs
            net.load_state_dict(states[E])
            net.eval()
            W0, b0 = last_layer(net)
            lam_score = simulate_lam(embed(net, Xm[tgt]), ym[tgt], W0, b0, cfg["lam_reps"])
            lam = max(LAM_GRID, key=lambda l: (lam_score[l], l))  # ties -> stronger pull
            sel = {"backbone": kind, "variant": var, "seed_used": int(seed),
                   "split": "lab.madrid_split(pm) defaults (frac .25, axis 0, high)",
                   "epoch_grid": list(grid),
                   "pseudo_target_f1_by_epoch": {str(k): float(v) for k, v in evals.items()},
                   "epochs": int(E), "lam_grid": list(LAM_GRID),
                   "pseudo_target_ft_f1_by_lam": {str(k): v for k, v in lam_score.items()},
                   "lam": float(lam),
                   "lam_sim": f"n per class in (5, 50) x {cfg['lam_reps']} reps on the pseudo-target, "
                              "argmax macro F1 on the rest",
                   "sec_per_full_epoch": sec, "n_src": int(len(src)), "n_pseudo_target": int(len(tgt))}
            sel_path.write_text(json.dumps(sel, indent=1))
            scores = {k: round(v, 4) for k, v in lam_score.items()}
            log(f"select {tag}: epochs={E} {sel['pseudo_target_f1_by_epoch']} lam={lam} {scores}")
        else:
            net, _, _, sec = train(kind, Xm[src], ym[src], Xa, sel["epochs"], seed,
                                   eval_at=(sel["epochs"],), max_rows=cfg["max_rows"],
                                   sched_epochs=max(sel["epoch_grid"]))
        Zm, Za = embed(net, Xm), embed(net, Xa)
        W0, b0 = last_layer(net)
        np.savez(path, Zm=Zm, Za=Za, W0=W0, b0=b0)
        embs.append((Zm, Za, W0, b0))
        log(f"backbone {tag} seed{seed}: {sec:.1f}s/full-epoch, {time.time() - t0:.0f}s")
    return sel, embs


def gru_check():
    """Smoke-only: the GRU network trains and embeds (the rest of its path is shared)."""
    rng = np.random.default_rng(0)
    X = rng.standard_normal((1024, 10, NY)).astype(np.float32)
    y = rng.integers(1, 5, 1024)
    net, _, _, sec = train("gru", X, y, None, 1, 0, eval_at=(1,))
    Z = embed(net, X[:300])
    assert Z.shape == (300, 64) and np.isfinite(Z).all()
    log(f"gru check ok: ~{sec * 61932 / 1024:.0f}s per Madrid-source epoch (extrapolated)")


# --------------------------------------------------------------------------- #
# main
# --------------------------------------------------------------------------- #
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--smoke", action="store_true")
    ap.add_argument("--only", default="")
    ap.add_argument("--extended", action="store_true",
                    help="also tcnn_dann_ry, gru_st, gru_ry (beyond the 70 min budget)")
    a = ap.parse_args()
    combos = list(DEFAULT_COMBOS) + (list(EXTENDED_COMBOS) if a.extended else [])
    if a.smoke:
        cfg = dict(tag="_SMOKE", suffix="_SMOKE", epoch_grid=(1, 2), max_rows=4096, lam_reps=1,
                   seeds=(0,), draws=[(5, 0), (200, 0)])
    else:
        cfg = dict(tag="", suffix="", epoch_grid=EPOCH_GRID, max_rows=None, lam_reps=5,
                   seeds=SEEDS, draws=list(lab.draws()))

    def mname(bb, var, head):
        return f"ts_{bb}_{var}_{head}{cfg['suffix']}"

    todo = {(bb, var): [h for h in HEADS if a.only in mname(bb, var, h)] for bb, var in combos}
    todo = {k: v for k, v in todo.items() if v}
    log(f"methods: {[mname(bb, var, h) for (bb, var), hs in todo.items() for h in hs]}")

    d = lab.load(official60=False)
    Tm = build_tensor("madrid", d["pm"])
    Ta = build_tensor("amsterdam", d["pa"])
    if a.smoke:
        gru_check()

    for (bb, var), heads in todo.items():
        need = [h for h in heads for (n, rep) in ([(0, 0)] if h == "zs" else cfg["draws"])
                if not lab.have_pred(mname(bb, var, h), n, rep)]
        if not need:
            log(f"{bb}_{var}: all predictions present, skip")
            continue
        Xm, Xa = make_inputs(Tm, var), make_inputs(Ta, var)
        log(f"inputs {var}: Madrid {Xm.shape} Amsterdam {Xa.shape}")
        sel, embs = backbone(bb, var, d, Xm, Xa, cfg)
        del Xm, Xa
        E, lam = sel["epochs"], sel["lam"]
        meta = {"family": "ts_nn", "backbone": bb, "variant": var, "seeds": list(cfg["seeds"]),
                "epochs": E, "lam": lam, "shot_weight": SHOT_WEIGHT, "selection": sel,
                "backbone_training_rows": "Madrid spatial source split only (lab.madrid_split)",
                "joint_input": "embeddings standardised with Madrid embedding mean/std",
                "input": f"10 ch x {NY} years (1984-2025), median of QA-valid obs, Blue<=BLUE_MAX",
                "augment": "band scale U(0.9,1.1), offset N(0,0.1)",
                "train": f"AdamW lr={LR} wd={WD} batch={BATCH} class-balanced CE, constant LR",
                "smoke": bool(a.smoke)}
        if "zs" in heads and not lab.have_pred(mname(bb, var, "zs"), 0, 0):
            P = np.mean([proba(Za, W0, b0) for _, Za, W0, b0 in embs], 0)
            lab.save_pred(mname(bb, var, "zs"), 0, 0, P, {**meta, "head": "zs"})
            log(f"{mname(bb, var, 'zs')} saved")
        scaled = []
        for Zm, Za, _, _ in embs:
            mu, sd = Zm.mean(0), Zm.std(0) + 1e-6
            scaled.append(((Zm - mu) / sd, (Za - mu) / sd))
        for n, rep in cfg["draws"]:
            hs = [h for h in heads if h != "zs" and not lab.have_pred(mname(bb, var, h), n, rep)]
            if not hs:
                continue
            t0 = time.time()
            shot = lab.shots_for(d, n, rep)
            ys = lab.shot_labels(d, shot)
            for h in hs:
                Ps = []
                for (Zm, Za, W0, b0), (Sm, Sa) in zip(embs, scaled):
                    if h == "ft":
                        Wf, bf = lab.finetune_last_layer(W0, b0, Za[shot], ys, lam)
                        Ps.append(proba(Za, Wf, bf))
                    else:
                        with warnings.catch_warnings():
                            warnings.simplefilter("ignore")
                            m = lab.joint_head(Sm, d["ym"], Sa[shot], ys, shot_weight=SHOT_WEIGHT)
                        Ps.append(m.predict_proba(Sa))
                lab.save_pred(mname(bb, var, h), n, rep, np.mean(Ps, 0), {**meta, "head": h})
            log(f"{bb}_{var} n={n} rep={rep} heads={hs} {time.time() - t0:.1f}s")
    log("done")


if __name__ == "__main__":
    main()