#!/usr/bin/env python3
"""Gnom trainer — KLD era (rewrite 2026-10-16).

    train_gnom.py --pool gnom/pool_v4.json --out gnom/models/v4.npz
                  [--variant trunk|modelbias|moeshared|rank|all] [--seed 7]
                  [--metric damage_kld] [--tail-feats]

Targets: the KLD family, 5 columns —
    damage_kld          mean KLD delta (primary, the bar column)
    damage_kld_max      tail / detonation
    damage_kld_median   body / skew
    damage_dp_rms       behavioral shift
    damage_same_top     top-1 agreement drop
All five arrive in one deterministic forward pass (new contract); the
loader reads any of them by name ('' as multi-target head sets).

DEAD METRICS (measured dead, kept only for archival pools):
    damage       PPL delta. Q5->Q4 std 0.77 BELOW the ~1.0 noise floor;
                 LOMO <= 0.06 across models (gnom/README, 2026-09/10).
    ssim_damage  structural. 4/4 structural metrics rank together and are
                 orthogonal to KLD/PPL (README SSIM section) — surrogate,
                 not discovery. No code path here does anything special
                 for it anyway.
Default --metric is and stays damage_kld.

Pool spec: [{"features": X.npz, "labels": L.jsonl, "model": tag,
             "base_bpw": 5.5, "drop_bpw": 4.5, "metric": "damage_kld"}, ...]
- null != 0. A null metric value (TERMINAL/censored record) is counted
  and skipped, never imputed (the guard this file has carried since
  2026-09-28, now explicit per-column).
- HOLE records live in *.failures.jsonl and are not read here at all.
- mapping check: every unit that carries a value must map to feature
  rows, tensor-complete. "mapping: U/U units, T/T tensors" is printed on
  every load; ANY miss prints details and exits 1 (--allow-partial-map
  downgrades the unit-with-no-match case to a warning). A pool that
  cannot prove its plumbing must not train.

Report is always: random split + LOMO (needs >= 2 models) + timp
baseline, each with its MDE at 80% power (alpha=0.05) next to it.
Bars: LOMO >= 0.45 (regression), rank LOMO-acc >= 0.60 (coin 0.50).
"""
import argparse
import copy
import json
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from tinynet import (init_params, train, forward, standardize, fit_mixed,
                     mixed_predict, moe_train_shared, moe_predict_shared,
                     backward_delta)

KLD_COLUMNS = ["damage_kld", "damage_kld_max", "damage_kld_median",
               "damage_dp_rms", "damage_same_top"]

# z_{0.975} + z_{0.80}: two-sided alpha 0.05, power 80%
Z_POWER = 1.959963984540054 + 0.8416212335729143  # ~2.8016


def mde_rho(n):
    """MDE for a single Spearman exam at 80% power (Fisher-z approx)."""
    return Z_POWER / np.sqrt(n - 3) if n > 3 else float("inf")


def mde_lomo(ns):
    """MDE of the MEAN LOMO rho across folds of sizes ns (z-space)."""
    ns = [n for n in ns if n > 3]
    if not ns:
        return float("inf")
    se = np.sqrt(sum(1.0 / (n - 3) for n in ns)) / len(ns)
    return Z_POWER * se


def mde_pair(n):
    """MDE for pair accuracy vs coin 0.50 at 80% power (n pairs)."""
    return Z_POWER * 0.5 / np.sqrt(n) if n > 0 else float("inf")


def build_pairs(y, models, max_pairs=30000, seed=7):
    """Within-model ordered pairs (a beats b): decisions, not values.
    Cross-model pairs are meaningless (damage scales differ per model).
    Returns (model, i, j, gap): RDO-style (2024/26 offline-RLHF) — HOW MUCH
    one is preferred matters, not just the order. Gap doubles as the
    resolution filter: pairs below measurement noise carry ~no weight."""
    rng = np.random.default_rng(seed)
    pairs = []
    for m in sorted(set(models.tolist() if hasattr(models, "tolist") else models)):
        idx = np.where(np.asarray(models) == m)[0]
        yy = y[idx]
        scale = float(np.std(yy) + 1e-12)
        cand = [(i, j, abs(float(yy[a] - y[j])) / scale)
                for a, i in enumerate(idx) for j in idx
                if yy[a] > y[j] + 1e-12]
        # (dead-zone tried 2026-10-03 and lost with the cap as a package:
        # uniform-over-all-pairs 0.575 stands as the best setting.)
        cand = [(i, j, g) for (i, j, g) in cand]
        if len(cand) > max_pairs:
            cand = [cand[k] for k in rng.choice(len(cand), max_pairs, replace=False)]
        pairs += [(m, i, j, g) for (i, j, g) in cand]
    return pairs


def pair_acc(scores, y, models, idx):
    """Fraction of within-model pairs ordered correctly."""
    idx = np.asarray(idx)
    tot, hit = 0, 0
    for m in set(np.asarray(models)[idx].tolist()):
        ii = idx[np.asarray(models)[idx] == m]
        yy = y[ii]
        ss = np.asarray(scores)[ii]
        for a in range(len(ii)):
            for b in range(len(ii)):
                if yy[a] > yy[b] + 1e-12:
                    tot += 1
                    if ss[a] > ss[b]:
                        hit += 1
    return hit / tot if tot else 0.5


def pair_count(y, models, idx):
    """Number of within-model ordered pairs in the eval set (for MDE)."""
    idx = np.asarray(idx)
    tot = 0
    for m in set(np.asarray(models)[idx].tolist()):
        ii = idx[np.asarray(models)[idx] == m]
        yy = y[ii]
        for a in range(len(ii)):
            for b in range(len(ii)):
                if yy[a] > yy[b] + 1e-12:
                    tot += 1
    return tot


def train_rank(ps, X, pairs_idx, lr=3e-3, epochs=1500, l2=1e-5, seed=7,
               activation="mish", log_every=1500):
    """RankNet on score differences (operator's System-1 idea): the net
    learns WHICH unit suffers more, never how much. Robust to label
    noise by construction — a flipped magnitude keeps its order."""
    rng = np.random.default_rng(seed)
    P = np.asarray([(a, b) for (_, a, b, _) in pairs_idx])
    W = np.asarray([g for (_, _, _, g) in pairs_idx], dtype=np.float64)
    W = W / W.sum()
    nval = max(4, int(len(P) * 0.2))
    vidx = rng.choice(len(P), nval, replace=False)
    mask = np.ones(len(P), bool)
    mask[vidx] = False
    trp, vap = P[mask], P[~mask]
    # NOTE (2026-10-03): RDO-style gap weighting tried here (W below) and
    # LOST (0.563 vs 0.575 unweighted). Why: with n=770 noisy labels the
    # big-gap pairs are outliers that mislead; the transfer boundary lives
    # in close pairs. Uniform stands; RDO left in code as an option.
    _USE_RDO = False  # verdict 2026-10-03: uniform 0.575 > huber 0.568 >
    # raw-RDO 0.563. Dead-zone killed the close pairs where the boundary
    # lives. Uniform stands, finally.
    Wtr, Wva = W[mask], W[~mask]
    if not _USE_RDO:
        Wtr = np.ones_like(Wtr) / len(Wtr)
    else:
        Wtr = np.minimum(Wtr, 1.0)  # Huber cap: nothing above 1 std counts extra
    m = [{k: np.zeros_like(v) for k, v in p.items()} for p in ps]
    v = [{k: np.zeros_like(v) for k, v in p.items()} for p in ps]
    best = (0.0, None, 0)
    for ep in range(1, epochs + 1):
        bi = rng.choice(len(trp), min(512, len(trp)), replace=False)
        A, B = trp[bi][:, 0], trp[bi][:, 1]
        fa, preca = forward(ps, X[A], activation)
        fb, precb = forward(ps, X[B], activation)
        d = (fa[-1].ravel() - fb[-1].ravel()) / 1.0
        sig = 1.0 / (1.0 + np.exp(np.clip(d, -30, 30)))
        wb = Wtr[bi] / Wtr[bi].sum()
        ga, gb = -sig * wb, +sig * wb
        grads = [{k: np.zeros_like(vv) for k, vv in p.items()} for p in ps]
        for acts, pre, g in ((fa, preca, ga), (fb, precb, gb)):
            delta = g[:, None]
            for i in reversed(range(len(ps))):
                grads[i]["W"] += acts[i].T @ delta + l2 * ps[i]["W"] / 2
                grads[i]["b"] += delta.sum(axis=0)
                if i > 0:
                    delta = backward_delta(delta @ ps[i]["W"].T, pre[i - 1],
                                           activation)
        for i in range(len(ps)):
            for k in ("W", "b"):
                g = grads[i][k]
                m[i][k] = 0.9 * m[i][k] + 0.1 * g
                v[i][k] = 0.95 * v[i][k] + 0.05 * g * g
                mh = m[i][k] / (1 - 0.9 ** ep)
                vh = v[i][k] / (1 - 0.95 ** ep)
                ps[i][k] -= lr * mh / (np.sqrt(vh) + 1e-8)
        pa = forward(ps, X[vap[:, 0]], activation)[0][-1].ravel()
        pb = forward(ps, X[vap[:, 1]], activation)[0][-1].ravel()
        acc = float(np.mean(pa > pb))
        if acc > best[0]:
            best = (acc, copy.deepcopy(ps), ep)
        if ep % log_every == 0 or ep == 1:
            print(f"  rank ep {ep}: val pair-acc={acc:.4f} "
                  f"(best {best[0]:.4f} @ {best[2]})", flush=True)
    if best[1] is not None:
        for i in range(len(ps)):
            ps[i] = best[1][i]
    return ps

# raw feature columns (tinynet.featurize order): 0-5 onehot, 6 layer,
# 7 logN, 8 timp_mean, 9 timp_max, 10 frag, 11 econc, 12-15 w-stats
SCALE_COLS = [7, 8, 9]


def spearman(a, b):
    ra = np.argsort(np.argsort(np.asarray(a, dtype=float))).astype(float)
    rb = np.argsort(np.argsort(np.asarray(b, dtype=float))).astype(float)
    ra -= ra.mean()
    rb -= rb.mean()
    d = np.sqrt((ra ** 2).sum() * (rb ** 2).sum())
    return float((ra * rb).sum() / d) if d > 0 else 0.0


def _unit_tag(names):
    # same rule as the allocator lookup (gnom.group_tag)
    # and predict_net: tissue@layer per blk tensor, "global" otherwise
    return "+".join(n.replace(".weight", "").split(".")[-1] + "@" +
                    n.split(".")[1] for n in names
                    if n.startswith("blk.")) or "global"


def load_pool_multi(path, metrics, force_metric=True, strict_map=True,
                    quiet=False, w_center=True, namefeats="off"):
    """Null-safe pooled loader for 1..k KLD-family targets.

    Rows are keyed on the PRIMARY metric (metrics[0]): a unit with a null
    primary is a censored observation — counted, skipped, never imputed.
    Other columns are NaN where the file does not carry them (old files
    predate the 5-column contract) or where the value is null.

    Returns (X, Y[n x k], models, gtags, stats). X columns: raw features
    (16) + base_bpw + drop_bpw [+ caller-appended extras].
    With w_center (default), the 4 weight-stat columns are centered on
    each model's own median, so their raw scale cannot leak across
    models — the exam is scale-invariant, the columns are not.
    """
    specs = json.load(open(path))
    Xs, Ys, models, gtags = [], [], [], []
    stats = {"units": 0, "mapped": 0, "tensors": 0, "missing": 0,
             "censored": 0, "files": 0, "n_feat_raw": 0,
             "missing_units": [], "files_detail": []}
    for sp in specs:
        primary = metrics[0] if force_metric else sp.get("metric", metrics[0])
        keys = [primary] + [m for m in metrics[1:]]
        z = np.load(sp["features"])
        X0 = np.asarray(z["X"], dtype=np.float64)
        stats["n_feat_raw"] = X0.shape[1]
        names = [str(t) for t in z["names"]]
        tmu = float(z["timp_mean"][0]) if "timp_mean" in z else 0.0
        tsd = float(z["timp_std"][0]) if "timp_std" in z else 1.0
        lmed = float(z["logN_med"][0]) if "logN_med" in z else 0.0
        row = {n: i for i, n in enumerate(names)}
        wmed = (np.median(X0[:, 12:16], axis=0)
                if w_center and X0.shape[1] >= 16 else None)
        base, rows, cens = None, [], 0
        with open(sp["labels"]) as f:
            for line in f:
                d = json.loads(line)
                if d["unit"] == "BASELINE":
                    base = d.get("ppl")
                    continue
                v0 = d.get(keys[0])
                if v0 is None:
                    # TERMINAL/censored (null != 0): visible in the file,
                    # never imputed into training. HOLEs are not in this
                    # file at all (*.failures.jsonl).
                    cens += 1
                    continue
                rows.append((d["unit"], [t for t in d.get("tensors", [])],
                             [d.get(k) for k in keys]))
        # DEAD PATH (archival pools only): PPL damage is relative to the
        # file's BASELINE ppl. KLD columns are already deltas (kf = 1).
        kf = 1.0 if keys[0] == "damage_kld" else (1.0 / base if base else 1.0)
        n_ok = t_ok = t_bad = 0
        for ui, (unit, tensors, vals) in enumerate(rows):
            have = [t for t in tensors if t in row]
            missing = [t for t in tensors if t not in row]
            stats["units"] += 1
            stats["tensors"] += len(tensors)
            stats["missing"] += len(missing)
            if missing:
                stats["missing_units"].append((sp["labels"], unit, missing))
                if not quiet:
                    print(f"  MAPPING MISS: {unit} -> {missing[:3]}"
                          f"{'...' if len(missing) > 3 else ''}", flush=True)
            if not have:
                continue
            idx = [row[t] for t in have]
            A = X0[idx]
            if wmed is not None:
                A = A.copy()
                A[:, 12:16] -= wmed  # weight stats vs model's own median
            g = A.mean(axis=0)
            # scale-invariant: timp z-scored in-model, logN vs model median
            g[8] = (A[:, 8].mean() - tmu) / (tsd + 1e-12)
            g[9] = (A[:, 9].max() - tmu) / (tsd + 1e-12)
            g[7] = A[:, 7].mean() - lmed
            if namefeats == "off":
                extra = []
            else:
                import hashlib as _hl
                if namefeats == "hash":
                    key = _unit_tag(tensors)
                elif namefeats == "layerabs":
                    ls = [n.split(".")[1] for n in tensors
                          if n.startswith("blk.") and
                          n.split(".")[1].isdigit()]
                    key = f"layer@{min(map(int, ls))}" if ls else "global"
                else:
                    key = f"anon@{ui}"
                b = _hl.md5(key.encode()).digest()[0] % 16
                extra = [1.0 if i == b else 0.0 for i in range(16)]
            Xs.append(np.concatenate([g, extra,
                                      [sp["base_bpw"], sp["drop_bpw"]]]))
            y = [vals[0] * kf]
            for v in vals[1:]:
                y.append(np.nan if v is None else v)
            Ys.append(y)
            models.append(sp["model"])
            gtags.append(f'{sp["model"]}:{keys[0]}:{unit}')
            n_ok += 1
            t_ok += len(have)
            t_bad += len(missing)
        stats["mapped"] += n_ok
        stats["censored"] += cens
        stats["files"] += 1
        stats["files_detail"].append(
            (sp["labels"], sp["features"], n_ok, len(rows), cens, t_ok, t_bad))
    if not quiet:
        for lab, feat, n_ok2, n_rows, cens, t_ok2, t_bad2 in stats["files_detail"]:
            print(f"  {lab}: mapped {n_ok2}/{n_rows}"
                  f" (tensors {t_ok2}/{t_ok2 + t_bad2}, censored {cens})"
                  f" [features {os.path.basename(feat)}]", flush=True)
        print(f"mapping: {stats['mapped']}/{stats['units']} units, "
              f"{stats['tensors'] - stats['missing']}/{stats['tensors']} tensors, "
              f"censored={stats['censored']}", flush=True)
    if strict_map and stats["missing"]:
        print(f"MAPPING FAIL: {len(stats['missing_units'])} units carry "
              f"tensors absent from their feature files — refusing to train "
              f"(--allow-partial-map to override)", flush=True)
        sys.exit(1)
    if stats["mapped"] < stats["units"] and strict_map:
        print(f"MAPPING FAIL: {stats['units'] - stats['mapped']} units with "
              f"values mapped to ZERO features — refusing to train",
              flush=True)
        sys.exit(1)
    return (np.array(Xs), np.array(Ys), np.array(models), gtags, stats)


def load_pool(path, metric="damage_kld", force_metric=False, strict_map=True):
    """Single-target view over load_pool_multi (historic signature)."""
    X, Y, models, gtags, stats = load_pool_multi(
        path, [metric], force_metric=force_metric, strict_map=strict_map)
    return X, Y[:, 0], models, gtags


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--pool", required=True)
    ap.add_argument("--out", default="gnom/models/v4.npz")
    ap.add_argument("--variant",
                    choices=["trunk", "modelbias", "moeshared", "rank", "all"],
                    default="all")
    ap.add_argument("--seed", type=int, default=7)
    ap.add_argument("--epochs", type=int, default=1500)
    ap.add_argument("--hidden", default="8,4",
                    help="hidden layer widths, comma-separated "
                         "(default '8,4'; try '16,8' for 2x capacity)")
    ap.add_argument("--metric", default="damage_kld",
                    help="label key; default damage_kld. 'damage' lists as "
                         "PPL — DEAD (archival pools only), see docstring.")
    ap.add_argument("--tail-feats", action="store_true",
                    help="append 4 explicit drop-depth x tissue interaction "
                         "columns (tail probe; see README)")
    ap.add_argument("--target-z", action="store_true",
                    help="z-score the target within each pool row (file): "
                         "equal footing for cross-corpus scales; the exam "
                         "(per-fold Spearman) is invariant to it")
    ap.add_argument("--allow-partial-map", action="store_true",
                    help="downgrade tensor-level mapping misses to warnings")
    ap.add_argument("--no-w-center", action="store_true",
                    help="keep raw (uncentered) weight-stat columns")
    ap.add_argument("--mask-timp", action="store_true",
                    help="zero the imatrix-derived columns (timp_mean/_max, "
                         "econc): statics + weight stats only, no lens")
    ap.add_argument("--namefeats", choices=["off", "hash", "anon",
                                                 "layerabs"],
                    default="off",
                    help="append a 16-bucket onehot of unit identity "
                         "(hash: md5(tag); anon: md5(unit index in file); "
                         "layerabs: md5(absolute layer number) — same "
                         "capacity everywhere; curiosity probe: "
                         "memorization vs transfer)")
    args = ap.parse_args()

    strict = not args.allow_partial_map
    X, Y, models, gtags, stats = load_pool_multi(
        args.pool, [args.metric], force_metric=True, strict_map=strict,
        w_center=not args.no_w_center, namefeats=args.namefeats)
    if args.namefeats != "off":
        print(f"namefeats ON ({args.namefeats}): +16 unit-identity columns",
              flush=True)
    y = Y[:, 0]  # single-target view: 1D, exactly the historic engine input
    if args.mask_timp:
        X[:, [8, 9, 11]] = 0.0  # before tail feats AND standardize/export
        print("mask-timp ON: imatrix columns zeroed (no-lens trunk)",
              flush=True)
    if args.tail_feats:
        # explicit tail interactions: drop depth x {body, tail, conc, attn}
        d = X[:, -2]
        attn = X[:, 0:3].sum(axis=1)
        extra = np.stack([d * X[:, 8], d * X[:, 9], d * X[:, 11],
                          d * attn], axis=1)
        X = np.concatenate([X, extra], axis=1)
        print("tail-feats ON: +4 columns (d*timp_mean, d*timp_max, d*econc, "
              "d*attn)", flush=True)
    y_raw = y.copy()  # exams always score against the raw target
    if args.target_z:
        for m in sorted(set(models)):
            msk = models == m
            yy = y[msk]
            y[msk] = (yy - yy.mean()) / (yy.std() + 1e-12)
        print("target-z ON: target z-scored within each pool row", flush=True)
    n_feat_raw = stats["n_feat_raw"] + (16 if args.namefeats != "off"
                                               else 0)
    print(f"pooled: n={len(y)}, models={sorted(set(models))}", flush=True)
    uniq_m = sorted(set(models))
    midx = np.array([uniq_m.index(m) for m in models])
    hidden = [int(h) for h in args.hidden.split(",") if h.strip()]
    DIMS = [X.shape[1]] + hidden + [1]
    print(f"arch: {'-'.join(map(str, DIMS))}", flush=True)

    def run_variant(variant, tr_idx, va_idx, seed):
        if variant == "rank":
            prs = build_pairs(y[tr_idx], midx[tr_idx])
            ps = init_params(DIMS, seed=seed)
            Xs, _, _ = standardize(X)
            train_rank(ps, Xs, prs, seed=seed)
            return forward(ps, Xs, "mish")[0][-1].ravel()
        if variant == "moeshared":
            exp = moe_train_shared(X[tr_idx][:, :n_feat_raw], y[tr_idx],
                                   seed=seed, epochs=args.epochs)
            return moe_predict_shared(exp, X[:, :n_feat_raw])
        if variant == "modelbias":
            # honest subset fit: trunk never sees va labels; unseen models
            # get bias 0 (pure trunk generalization).
            Xtr, ytr = X[tr_idx], y[tr_idx]
            mtr = midx[tr_idx]
            ug = sorted(set(mtr.tolist()))
            r = {g: i for i, g in enumerate(ug)}
            gs = np.array([r[g] for g in mtr])
            mp = fit_mixed(Xtr, ytr, gs, len(ug), DIMS, seed=seed,
                           epochs=args.epochs, log_every=args.epochs)
            Xs = (X - mp["xmu"]) / mp["xsd"]
            base = forward(mp["trunk"], Xs, "mish")[0][-1].ravel() \
                * mp["tsd"] + mp["tmu"]
            b = np.zeros(len(X))
            known = np.array([r[m] if m in r else -1 for m in midx])
            kk = known >= 0
            b[kk] = mp["bias"][known[kk]]
            return (base + b) * mp["ysd"] + mp["ymu"]
        ps = init_params(DIMS, seed=seed)
        Xs, _, _ = standardize(X)
        w = np.abs(y[tr_idx]) + 0.05 * np.abs(y[tr_idx]).max()
        train(ps, Xs[tr_idx], y[tr_idx], epochs=args.epochs,
              log_every=args.epochs, noise_std=0.05, sample_w=w,
              activation="mish", dropout=0.0, seed=seed)
        mu, sd = float(y.mean()), float(y.std() + 1e-12)
        return forward(ps, Xs, "mish")[0][-1].ravel() * sd + mu

    variants = (["trunk", "modelbias", "moeshared", "rank"]
                if args.variant == "all" else [args.variant])
    rng = np.random.default_rng(args.seed)
    # random split (group-disjoint)
    uniq_g = sorted(set(gtags))
    gperm = rng.permutation(len(uniq_g))
    gtr = set(gperm[:max(2, len(uniq_g) * 4 // 5)])
    gidx = np.array([uniq_g.index(t) for t in gtags])
    tr = np.array([i for i in range(len(X)) if gidx[i] in gtr])
    va = np.array([i for i in range(len(X)) if gidx[i] not in gtr])
    print(f"split: train {len(tr)} / heldout {len(va)}", flush=True)
    results = {}
    for v in variants:
        p = run_variant(v, tr, va, args.seed)
        if v == "rank":
            # decisions are scored by pair accuracy, not correlation:
            # coin is 0.50, bar is 0.60.
            acc_r = pair_acc(p, y_raw, midx, va)
            np_r = pair_count(y_raw, midx, va)
            acc_lomo, np_lomo = [], []
            if len(uniq_m) > 1:
                for m in uniq_m:
                    trm = np.where(models != m)[0]
                    vam = np.where(models == m)[0]
                    if len(vam) < 5 or len(trm) < 5:
                        continue
                    pm = run_variant(v, trm, vam, args.seed)
                    acc_lomo.append(pair_acc(pm, y_raw, midx, vam))
                    np_lomo.append(pair_count(y_raw, midx, vam))
            lomo = float(np.mean(acc_lomo)) if acc_lomo else 0.5
            mde_r = mde_pair(np_r)
            mde_l = mde_pair(np.mean(np_lomo)) if np_lomo else float("inf")
            print(f"{v}: random-acc={acc_r:.3f} (MDE={mde_r:.3f}, "
                  f"n={np_r}) LOMO-acc={lomo:.3f} (MDE~{mde_l:.3f}) "
                  f"(per-model {[f'{s:.2f}' for s in acc_lomo]})", flush=True)
            results[v] = {"random": acc_r, "lomo": lomo}
            continue
        sp_r = spearman(p[va], y_raw[va])
        # LOMO needs >=2 models; single-metric pools report random only.
        sp_lomo, fold_n, timp_lomo = [], [], []
        if len(uniq_m) > 1:
            for m in uniq_m:
                trm = np.where(models != m)[0]
                vam = np.where(models == m)[0]
                if len(vam) < 5 or len(trm) < 5:
                    continue
                pm = run_variant(v, trm, vam, args.seed)
                sp_lomo.append(spearman(pm[vam], y_raw[vam]))
                fold_n.append(len(vam))
                timp_lomo.append(spearman(X[vam, 8], y_raw[vam]))
        lomo = float(np.mean(sp_lomo)) if sp_lomo else 0.0
        base = spearman(X[va, 8], y_raw[va])
        tb_lomo = float(np.mean(timp_lomo)) if timp_lomo else None
        mde_r = mde_rho(len(va))
        mde_l = mde_lomo(fold_n)
        tb_str = (f"timp_base={base:.3f} timp_LOMO="
                  f"{tb_lomo:.3f}" if tb_lomo is not None
                  else f"timp_base={base:.3f} timp_LOMO=n/a")
        print(f"{v}: random={sp_r:.3f} (MDE={mde_r:.3f}, n={len(va)}) "
              f"LOMO={lomo:.3f} (MDE~{mde_l:.3f}, folds={fold_n}) "
              f"{tb_str}", flush=True)
        print(f"  per-model {[f'{s:.2f}' for s in sp_lomo]}"
              + (f" timp {[f'{s:.2f}' for s in timp_lomo]}"
                 if timp_lomo else ""), flush=True)
        results[v] = {"random": sp_r, "lomo": lomo}
    best_c = max([v for v in results if v != "rank"],
                 key=lambda v: results[v]["lomo"], default=None)
    if best_c:
        print(f"BEST-corr: {best_c} LOMO={results[best_c]['lomo']:.3f} "
              f"(bar 0.45: {'PASS' if results[best_c]['lomo'] >= 0.45 else 'FAIL'})",
              flush=True)
    if "rank" in results:
        print(f"BEST-rank: rank LOMO-acc={results['rank']['lomo']:.3f} "
              f"(bar 0.60, coin 0.50: "
              f"{'PASS' if results['rank']['lomo'] >= 0.60 else 'FAIL'})",
              flush=True)
    # Export: trunk trained on FULL data (moeshared/modelbias export TODO).
    # Whole gnom: ~220 floats ≈ 2KB in one .npz (precompute per model into
    # netpred.json; no live inference needed in the queue).
    # predict_net rebuilds hash-identity columns from unit tags, but
    # anon indices depend on labels-file order: anon trunks are
    # metrics-only, skipped here rather than shipping an unloadable model.
    if args.namefeats == "anon":
        print("anon trunk: export skipped (predict cannot rebuild "
              "file-order indices)", flush=True)
        return
    ps = init_params(DIMS, seed=args.seed)
    Xs, xmu, xsd = standardize(X)
    ymu, ysd = float(y.mean()), float(y.std() + 1e-12)
    w = np.abs(y) + 0.05 * np.abs(y).max()
    train(ps, Xs, (y - ymu) / ysd, epochs=args.epochs,
          log_every=args.epochs, noise_std=0.05, sample_w=w,
          activation="mish", dropout=0.0, seed=args.seed)
    _os = os
    _os.makedirs(_os.path.dirname(_os.path.abspath(args.out)) or ".",
                 exist_ok=True)
    np.savez(args.out, xmu=xmu, xsd=xsd, ymu=np.array(ymu), ysd=np.array(ysd),
             dims=np.array(DIMS), metric=np.array(args.metric),
             **{f"p{i}{k}": v for i, p in enumerate(ps) for k, v in p.items()})
    print(f"saved trunk {args.out}", flush=True)


if __name__ == "__main__":
    main()