"""
evaluate.py - leakage-free predictability analysis: which workbook quantities can be predicted
from hive sound, and which cannot?

Protocol, applied identically to every (target, feature family) pair
  * colony-grouped cross-validation (5 folds x 3 repeats, every recording of a colony in the
    same fold) -> out-of-fold (OOF) predictions; repeats are averaged
  * metrics: ROC-AUC and average precision (binary); Spearman rho, Pearson r, R2 (regression)
  * uncertainty: colony-cluster bootstrap 95 % confidence interval (colonies resampled; the
    fraction of resamples without both classes is recorded and the CI withheld if it exceeds 5 %)
  * significance: month-aligned colony permutation test - within each country the colony
    identities are permuted and every recording receives the label that the donor colony had in
    the SAME calendar month (recordings whose donor has no recording in that month are left out
    of that permutation). Country, calendar month, the October 2023 batch and the within-colony
    label structure all survive the permutation; p-values therefore test for acoustic signal
    beyond site, season and batch.
  * confound control:
      - metadata-only baseline (site, two seasonal harmonics, year, October 2023 batch) and the
        paired bootstrap of the gain of "metadata + audio" over "metadata"
      - within-country evaluation with its own metadata baseline and paired gain
      - cross-site transfer (train on one country, test on the other)
      - every target re-run without the October 2023 single-visit batch
  * evaluations are gated on the number of positive AND negative colonies (binary) or rows and
    colonies (regression); anything below the gate is reported as "not estimable"
  * models are fixed a priori (no per-target model selection):
      linear  : impute -> standardise -> PCA(64, only if more than 64 inputs; capped by the
                training-set size) -> L2 logistic regression (C = 0.1, balanced) / ridge (alpha = 10)
      hgb     : histogram gradient boosting (balanced class weights) on raw features
      ensemble: rank-average of the OOF predictions of the ast, clap and hc families (primary)
  * verdict rule (pre-specified):
      level 0  no evidence: the primary CI includes chance (the upper bound is reported)
      level 1  signal present but explained by site/season/batch: CI excludes chance, but the
               permutation p >= 0.05 or the paired gain over metadata is not > 0
      level 2  predictable beyond site/season/batch: CI excludes chance, permutation p < 0.05
               AND paired gain over metadata > 0 (Holm-adjusted p across targets also reported)

Outputs (artifacts/results/):
  results.json, results_long.csv, oof/<target>.npz, permutation.json (with --n-perm > 0)
"""
import argparse
import json
import os
import sys
import time
import warnings

import numpy as np
import pandas as pd
from scipy.stats import spearmanr, pearsonr, rankdata
from sklearn.decomposition import PCA
from sklearn.ensemble import HistGradientBoostingClassifier, HistGradientBoostingRegressor
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression, Ridge
from sklearn.metrics import roc_auc_score, average_precision_score
from sklearn.model_selection import StratifiedGroupKFold
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import data as bd  # noqa: E402

warnings.filterwarnings("ignore")
FEAT_DIR = os.path.join(HERE, "artifacts", "features")
RES_DIR = os.path.join(HERE, "artifacts", "results")
PCA_K = 64
LIN_C = 0.1
RIDGE_ALPHA = 10.0
N_SPLITS = 5
N_REPEATS = 3
N_BOOT = 1000
MIN_POS_COLONIES = 8        # binary: at least this many positive AND negative colonies
MIN_ROWS_REG = 40           # regression: at least this many rows ...
MIN_COLONIES_REG = 10       # ... from at least this many colonies
MAX_SKIPPED_BOOT = 0.05     # withhold a CI if more resamples than this lacked both classes
FAMILY_MODEL = {            # a priori model per family
    "meta": "lin", "qc": "hgb", "gem": "lin", "hc": "hgb", "ast": "lin", "clap": "lin",
    "ast+hc": "lin", "all_audio": "lin",
    "meta+hc": "hgb", "meta+ast": "lin", "meta+all_audio": "lin",
    "lab+meta": "hgb", "lab+meta+all_audio": "lin",
}
AUDIO_FAMILIES = ["hc", "ast", "clap", "ast+hc", "all_audio", "ensemble"]
ENSEMBLE_MEMBERS = ["ast", "clap", "hc"]
PERM_FAMILIES = ["meta", "hc", "ast", "clap", "all_audio", "meta+all_audio"]
SENS_FAMILIES = ["meta", "hc", "ast", "clap", "all_audio", "meta+all_audio"]
OCT2023_BATCH = "USA_2023_10"


# --------------------------------------------------------------------------- features
def load_features(df, feat_dir=FEAT_DIR, exclude=(), verbose=True):
    """Return (df_aligned, {family: X}, names, coverage) restricted to recordings present in every table."""
    key = df["sound_file"].str.lower().values
    fam, names = {}, {}

    def _csv(name, prefix=None):
        p = os.path.join(feat_dir, name)
        if not os.path.exists(p):
            return None, None
        t = pd.read_csv(p)
        t["file"] = t["file"].astype(str).str.lower()
        t = t.drop_duplicates("file").set_index("file")
        cols = [c for c in t.columns if pd.api.types.is_numeric_dtype(t[c])
                and (prefix is None or c.startswith(prefix))]
        return t[cols], cols

    def _npz(name):
        p = os.path.join(feat_dir, name)
        if not os.path.exists(p):
            return None
        z = np.load(p, allow_pickle=False)
        files = [str(f).lower() for f in z["files"]]
        return pd.DataFrame(z["X"].astype(np.float32), index=files)

    tabs = {}
    t, c = _csv("handcrafted.csv"); tabs["hc"] = t; names["hc"] = c
    t, c = _csv("qc.csv"); tabs["qc"] = t; names["qc"] = c
    t, c = _csv("gemini_semantic.csv", prefix="gem_"); tabs["gem"] = t; names["gem"] = c
    tabs["ast"] = _npz("ast_pooled.npz")
    tabs["clap"] = _npz("clap_pooled.npz")
    tabs = {k: v for k, v in tabs.items() if v is not None and k not in set(exclude)}
    coverage = {k: int(pd.Index(key).isin(v.index).sum()) for k, v in tabs.items()}
    common = None
    for k, v in tabs.items():
        s = set(v.index)
        common = s if common is None else common & s
    keep = np.array([k in common for k in key])
    dropped = df.loc[~keep]
    if verbose:
        print("feature coverage (rows of the workbook found per table): " + ", ".join(f"{k}={v}" for k, v in coverage.items()), flush=True)
        if len(dropped):
            print(f"WARNING: {len(dropped)} recordings dropped because a feature table lacks them; colonies affected: "
                  f"{sorted(set(dropped['colony']))[:20]}", flush=True)
    df = df[keep].reset_index(drop=True)
    key = key[keep]
    for k, v in tabs.items():
        fam[k] = v.reindex(key).values.astype(np.float32)
    fam["meta"], names["meta"] = bd.metadata_matrix(df)
    lab_cols = [c for c in df.columns if c not in bd.ID_COLUMNS and c not in ("y03", "y46")
                and pd.api.types.is_numeric_dtype(df[c])]
    if lab_cols:
        fam["lab"] = np.log1p(np.clip(df[lab_cols].astype(float).fillna(0).values, 0, None)).astype(np.float32)
        names["lab"] = [f"log1p_{c}" for c in lab_cols]
    if "ast" in fam and "hc" in fam:
        fam["ast+hc"] = np.hstack([fam["ast"], fam["hc"]])
    audio = [fam[k] for k in ("ast", "clap", "hc") if k in fam]
    if audio:
        fam["all_audio"] = np.hstack(audio)
    if "hc" in fam:
        fam["meta+hc"] = np.hstack([fam["meta"], fam["hc"]])
    if "ast" in fam:
        fam["meta+ast"] = np.hstack([fam["meta"], fam["ast"]])
    if "all_audio" in fam:
        fam["meta+all_audio"] = np.hstack([fam["meta"], fam["all_audio"]])
        if "lab" in fam:
            fam["lab+meta"] = np.hstack([fam["lab"], fam["meta"]])
            fam["lab+meta+all_audio"] = np.hstack([fam["lab"], fam["meta"], fam["all_audio"]])
    return df, fam, names, dict(coverage=coverage, n_dropped=int(len(dropped)))


# --------------------------------------------------------------------------- models
class AdaptivePCA(PCA):
    """PCA whose number of components is capped at fit time by the training-set size."""
    def fit(self, X, y=None):
        k = int(min(self.n_components, X.shape[1], max(1, X.shape[0] - 1)))
        self.set_params(n_components=k)
        return super().fit(X, y)

    def fit_transform(self, X, y=None):
        self.fit(X, y)
        return self.transform(X)


def make_model(kind, model, n_features):
    if model == "lin":
        steps = [("imp", SimpleImputer(strategy="median")), ("sc", StandardScaler())]
        if n_features > PCA_K:
            steps.append(("pca", AdaptivePCA(n_components=PCA_K, random_state=0)))
        if kind == "binary":
            steps.append(("clf", LogisticRegression(C=LIN_C, class_weight="balanced", max_iter=5000)))
        else:
            steps.append(("reg", Ridge(alpha=RIDGE_ALPHA)))
        return Pipeline(steps)
    if model == "hgb":
        common = dict(max_iter=200, learning_rate=0.06, max_leaf_nodes=15, min_samples_leaf=20,
                      l2_regularization=1.0, random_state=0)
        if kind == "binary":
            return HistGradientBoostingClassifier(class_weight="balanced", **common)
        return HistGradientBoostingRegressor(**common)
    raise ValueError(model)


def _predict(m, kind, X):
    return m.predict_proba(X)[:, 1] if kind == "binary" else m.predict(X)


# --------------------------------------------------------------------------- gates
def colony_counts(kind, y, groups):
    if kind == "binary":
        pos = set(groups[y == 1]); neg = set(groups[y == 0])
        return dict(n_pos_colonies=len(pos), n_neg_colonies=len(neg), n_colonies=len(set(groups)))
    return dict(n_colonies=len(set(groups)))


def enough(kind, y, groups):
    """(ok, reason) - is this subset large enough to evaluate?"""
    cc = colony_counts(kind, y, groups)
    if kind == "binary":
        if cc["n_pos_colonies"] < MIN_POS_COLONIES or cc["n_neg_colonies"] < MIN_POS_COLONIES:
            return False, f"not estimable: {cc['n_pos_colonies']} positive / {cc['n_neg_colonies']} negative colonies (need {MIN_POS_COLONIES} each)"
    else:
        if len(y) < MIN_ROWS_REG or cc["n_colonies"] < MIN_COLONIES_REG:
            return False, f"not estimable: {len(y)} rows from {cc['n_colonies']} colonies (need {MIN_ROWS_REG} rows, {MIN_COLONIES_REG} colonies)"
    if cc["n_colonies"] < N_SPLITS:
        return False, "not estimable: fewer colonies than folds"
    return True, ""


# --------------------------------------------------------------------------- CV
def group_splits(y, groups, kind, seed, n_splits=N_SPLITS):
    """Colony-grouped folds. Stratified on the label for binary targets."""
    idx = np.arange(len(y))
    if kind == "binary":
        sgk = StratifiedGroupKFold(n_splits=n_splits, shuffle=True, random_state=seed)
        return list(sgk.split(idx, y.astype(int), groups))
    rng = np.random.RandomState(seed)
    ug = np.array(sorted(set(groups)))
    rng.shuffle(ug)
    return [(idx[~np.isin(groups, f)], idx[np.isin(groups, f)]) for f in np.array_split(ug, n_splits)]


def oof_predict(X, y, groups, kind, model, seeds=range(N_REPEATS), n_splits=N_SPLITS):
    """Average OOF prediction over repeats (each repeat covers every row once)."""
    P = np.zeros((len(seeds), len(y)), dtype=np.float64)
    for r, seed in enumerate(seeds):
        for tr, te in group_splits(y, groups, kind, seed, n_splits):
            m = make_model(kind, model, X.shape[1]).fit(X[tr], y[tr])
            P[r, te] = _predict(m, kind, X[te])
    return P.mean(axis=0)


def metric(kind, y, p):
    if kind == "binary":
        if len(np.unique(y)) < 2:
            return dict(auc=np.nan, ap=np.nan)
        return dict(auc=float(roc_auc_score(y, p)), ap=float(average_precision_score(y, p)))
    if np.std(p) == 0 or np.std(y) == 0:
        return dict(spearman=0.0, pearson=0.0, r2=float(1 - np.sum((y - p) ** 2) / np.sum((y - y.mean()) ** 2)))
    return dict(spearman=float(spearmanr(y, p)[0]), pearson=float(pearsonr(y, p)[0]),
                r2=float(1 - np.sum((y - p) ** 2) / np.sum((y - y.mean()) ** 2)))


def primary_name(kind):
    return "auc" if kind == "binary" else "spearman"


def chance_level(kind):
    return 0.5 if kind == "binary" else 0.0


def cluster_bootstrap(kind, y, preds, groups, n_boot=N_BOOT, seed=0):
    """preds: dict name -> OOF vector. Returns (ci {name: (lo, hi)}, paired deltas, skipped fraction)."""
    rng = np.random.RandomState(seed)
    ug = np.array(sorted(set(groups)))
    gi = {g: np.where(groups == g)[0] for g in ug}
    stat = {k: [] for k in preds}
    skipped = 0
    for _ in range(n_boot):
        pick = rng.choice(ug, size=len(ug), replace=True)
        idx = np.concatenate([gi[g] for g in pick])
        yb = y[idx]
        if kind == "binary" and len(np.unique(yb)) < 2:
            skipped += 1
            continue
        for k, p in preds.items():
            stat[k].append(metric(kind, yb, p[idx])[primary_name(kind)])
    frac = skipped / max(1, n_boot)
    if frac > MAX_SKIPPED_BOOT:
        return {k: (np.nan, np.nan) for k in preds}, {}, frac
    ci = {k: (float(np.nanpercentile(v, 2.5)), float(np.nanpercentile(v, 97.5))) if len(v) else (np.nan, np.nan)
          for k, v in stat.items()}
    delta = {}
    pairs = [(k, "meta") for k in stat if k.startswith("meta+") and "meta" in stat]
    if "lab+meta+all_audio" in stat and "lab+meta" in stat:
        pairs.append(("lab+meta+all_audio", "lab+meta"))
    for k, base_k in pairs:
        base = np.array(stat[base_k]); v = stat[k]
        if len(v) == len(base) and len(v):
            d = np.array(v) - base
            delta[k] = (float(np.nanpercentile(d, 2.5)), float(np.nanpercentile(d, 97.5)), float(np.nanmean(d)))
    return ci, delta, frac


def rank_average(vectors):
    return np.mean([rankdata(v) / len(v) for v in vectors], axis=0)


# --------------------------------------------------------------------------- permutation
def colony_permute(y, groups, country, t_index, rng):
    """Month-aligned colony permutation within country.

    Colony identities are permuted within each country; every recording of the recipient colony
    receives the label the donor colony had in the same calendar month. Recordings whose donor has
    no recording in that month get NaN (excluded from that permutation). Site, month, batch and
    the within-colony label structure are preserved; only the link to the colony's audio is broken.
    """
    yp = np.full(len(y), np.nan)
    lookup = {(g, t): v for g, t, v in zip(groups, t_index, y)}
    for c in sorted(set(country)):
        cols = np.array(sorted(set(groups[country == c])))
        donors = dict(zip(cols, cols[rng.permutation(len(cols))]))
        for i in np.where(country == c)[0]:
            yp[i] = lookup.get((donors[groups[i]], t_index[i]), np.nan)
    return yp


def _fit_families(fams, fam, mask, y, groups, kind, seed):
    pp = {}
    for f in fams:
        pp[f] = oof_predict(fam[f][mask], y, groups, kind, FAMILY_MODEL[f], seeds=[seed])
    mem = [f for f in ENSEMBLE_MEMBERS if f in pp]
    if len(mem) >= 2:
        pp["ensemble"] = rank_average([pp[f] for f in mem])
    return pp


def permutation_target(t, fam, df, families, n_perm, seed=0):
    kind, y_all = t["kind"], t["y"]
    mask = ~np.isnan(y_all)
    y = y_all[mask]
    groups = df["colony"].values[mask]
    country = df["country"].values[mask]
    t_index = df["t_index"].astype(int).values[mask]
    ok, why = enough(kind, y, groups)
    if not ok:
        return {"skipped": why}
    rng = np.random.RandomState(seed)
    pn = primary_name(kind)
    fams = [f for f in families if f in fam]
    obs = {k: metric(kind, y, p)[pn] for k, p in _fit_families(fams, fam, mask, y, groups, kind, 0).items()}
    null = {k: [] for k in obs}
    n_used = []
    for b in range(n_perm):
        yp = colony_permute(y, groups, country, t_index, rng)
        keep = ~np.isnan(yp)
        yb, gb = yp[keep], groups[keep]
        if not enough(kind, yb, gb)[0]:
            continue
        sub = np.where(mask)[0][keep]
        m2 = np.zeros(len(y_all), dtype=bool); m2[sub] = True
        pp = _fit_families(fams, fam, m2, yb, gb, kind, b + 100)
        for k, p in pp.items():
            null[k].append(metric(kind, yb, p)[pn])
        n_used.append(int(keep.sum()))
    out = {}
    for k, o in obs.items():
        nl = np.array(null[k])
        if len(nl) == 0:
            continue
        out[k] = dict(observed_single_repeat=float(o), null_mean=float(nl.mean()), null_sd=float(nl.std()),
                      null_p95=float(np.percentile(nl, 95)), p_value=float((np.sum(nl >= o) + 1) / (len(nl) + 1)),
                      n_perm=int(len(nl)), mean_rows_per_perm=float(np.mean(n_used)) if n_used else None)
    return out


# --------------------------------------------------------------------------- per-target
def _eval_block(fam, mask, y, groups, kind, families, seeds, n_boot):
    preds = {}
    for f in families:
        if f in fam:
            preds[f] = oof_predict(fam[f][mask], y, groups, kind, FAMILY_MODEL[f], seeds)
    members = [preds[k] for k in ENSEMBLE_MEMBERS if k in preds]
    if len(members) >= 2:
        preds["ensemble"] = rank_average(members)
    point = {k: metric(kind, y, p) for k, p in preds.items()}
    ci, delta, frac = cluster_bootstrap(kind, y, preds, groups, n_boot=n_boot)
    fams = {k: dict(**point[k], ci=list(ci[k])) for k in preds}
    deltas = {k: dict(lo=v[0], hi=v[1], mean=v[2]) for k, v in delta.items()}
    return preds, fams, deltas, frac


def evaluate_target(t, fam, df, families, n_repeats=N_REPEATS, n_boot=N_BOOT, verbose=True):
    kind, y_all = t["kind"], t["y"]
    mask = ~np.isnan(y_all)
    y = y_all[mask]
    groups = df["colony"].values[mask]
    country = df["country"].values[mask]
    batch = df["recording_batch"].values[mask]
    res = dict(name=t["name"], kind=kind, description=t["description"], source=t["source"],
               n=int(mask.sum()), **colony_counts(kind, y, groups))
    if kind == "binary":
        res["n_pos"] = int(y.sum()); res["prevalence"] = float(y.mean())
        res["n_pos_rows_oct2023"] = int(((batch == OCT2023_BATCH) & (y == 1)).sum())
    ok, why = enough(kind, y, groups)
    if not ok:
        res["skipped"] = why
        return res, None
    seeds = list(range(n_repeats))
    t0 = time.time()
    preds, fams, deltas, frac = _eval_block(fam, mask, y, groups, kind, families, seeds, n_boot)
    res["families"] = fams
    res["delta_over_meta"] = deltas
    res["bootstrap_skipped_fraction"] = frac

    # within-country: audio, metadata and their paired gain, inside one country
    res["within_country"] = {}
    for c in sorted(set(country)):
        m = country == c
        okc, whyc = enough(kind, y[m], groups[m])
        if not okc:
            res["within_country"][c] = dict(skipped=whyc, n=int(m.sum()))
            continue
        idx = np.where(mask)[0][m]
        m2 = np.zeros(len(y_all), dtype=bool); m2[idx] = True
        wf = [f for f in ("meta", "hc", "ast", "clap", "all_audio", "meta+all_audio") if f in fam]
        p_c, f_c, d_c, fr_c = _eval_block(fam, m2, y[m], groups[m], kind, wf, seeds, max(200, n_boot // 2))
        res["within_country"][c] = dict(n=int(m.sum()), **colony_counts(kind, y[m], groups[m]),
                                        families=f_c, delta_over_meta=d_c, bootstrap_skipped_fraction=fr_c)

    # cross-site transfer
    res["transfer"] = {}
    for f in ("ast", "hc", "all_audio"):
        if f not in fam:
            continue
        X = fam[f][mask]
        for src, dst in (("USA", "NZ"), ("NZ", "USA")):
            tr = country == src; te = country == dst
            if not enough(kind, y[tr], groups[tr])[0] or not enough(kind, y[te], groups[te])[0]:
                continue
            m = make_model(kind, FAMILY_MODEL[f], X.shape[1]).fit(X[tr], y[tr])
            p = _predict(m, kind, X[te])
            cic, _, frt = cluster_bootstrap(kind, y[te], {f: p}, groups[te], n_boot=max(200, n_boot // 2))
            res["transfer"].setdefault(f, {})[f"{src}->{dst}"] = dict(n_test=int(te.sum()), **metric(kind, y[te], p), ci=list(cic[f]),
                                                                      bootstrap_skipped_fraction=frt)

    # sensitivity: every target without the October 2023 single-visit batch
    keep = batch != OCT2023_BATCH
    if keep.sum() < mask.sum():
        oks, whys = enough(kind, y[keep], groups[keep])
        if oks:
            idx = np.where(mask)[0][keep]
            m2 = np.zeros(len(y_all), dtype=bool); m2[idx] = True
            sf = [f for f in SENS_FAMILIES if f in fam]
            p_s, f_s, d_s, fr_s = _eval_block(fam, m2, y[keep], groups[keep], kind, sf, seeds, max(200, n_boot // 2))
            res["without_oct2023"] = dict(n=int(keep.sum()), **colony_counts(kind, y[keep], groups[keep]),
                                          n_pos=int(y[keep].sum()) if kind == "binary" else None,
                                          families=f_s, delta_over_meta=d_s, bootstrap_skipped_fraction=fr_s)
        else:
            res["without_oct2023"] = dict(skipped=whys, n=int(keep.sum()))
    res["seconds"] = round(time.time() - t0, 1)
    if verbose:
        pn = primary_name(kind)
        best = max(((k, v[pn]) for k, v in fams.items() if k in AUDIO_FAMILIES), key=lambda kv: kv[1], default=("-", np.nan))
        print(f"  {t['name']:34s} {kind:10s} n={res['n']:4d} meta {fams.get('meta', {}).get(pn, np.nan):.3f} "
              f"ensemble {fams.get('ensemble', {}).get(pn, np.nan):.3f} best-audio {best[0]} {best[1]:.3f}  [{res['seconds']}s]", flush=True)
    oof = dict(y=y, groups=groups, country=country, **{f"pred_{k}": v for k, v in preds.items()})
    return res, oof


# --------------------------------------------------------------------------- verdicts
def strength(kind, v):
    if kind == "binary":
        return "strong" if v >= 0.75 else "moderate" if v >= 0.65 else "weak"
    return "strong" if abs(v) >= 0.4 else "moderate" if abs(v) >= 0.2 else "weak"


def verdict(res, perm=None):
    """Pre-specified decision rule; primary family = ensemble (fallback: all_audio)."""
    if "skipped" in res:
        return dict(level=-1, label="not evaluated", reason=res["skipped"])
    kind = res["kind"]
    pn = primary_name(kind)
    fams = res["families"]
    prim = "ensemble" if "ensemble" in fams else "all_audio"
    if prim not in fams:
        return dict(level=-1, label="not evaluated", reason="no audio family")
    v = fams[prim][pn]; lo, hi = fams[prim]["ci"]
    chance = chance_level(kind)
    out = dict(primary=prim, value=v, ci=[lo, hi])
    if lo is None or np.isnan(lo):
        out.update(level=-1, label="not evaluated", reason="bootstrap CI withheld (too few positive colonies)")
        return out
    if kind == "binary":
        sig = lo > chance
        bound = f"AUC at most {hi:.2f} with 95% confidence"
    else:
        sig = lo > 0 or hi < 0
        bound = f"|Spearman| at most {max(abs(lo), abs(hi)):.2f} with 95% confidence"
    p = None
    if perm and prim in perm and "p_value" in perm[prim]:
        p = perm[prim]["p_value"]
    out["p_perm"] = p
    d = res.get("delta_over_meta", {}).get("meta+all_audio")
    out["gain_over_meta"] = d
    beyond_meta = d is not None and d["lo"] > 0
    out["beyond_meta"] = beyond_meta
    wo = (res.get("without_oct2023") or {}).get("families", {}).get(prim)
    if wo:
        out["without_oct2023"] = dict(value=wo[pn], ci=wo["ci"])
    if not sig:
        out.update(level=0, label="no evidence of predictability from sound", upper_bound=bound,
                   reason=f"pooled colony-grouped performance not distinguishable from chance ({bound})")
        return out
    perm_ok = (p is not None and p < 0.05)
    if perm_ok and beyond_meta:
        s = strength(kind, v)
        out.update(level=2, label=f"predictable beyond site/season/batch ({s})", strength=s,
                   reason="permutation p < 0.05 with month-aligned colony permutation and paired gain over metadata > 0")
        return out
    why = []
    if p is None:
        why.append("no permutation test run")
    elif not perm_ok:
        why.append(f"month-aligned permutation p = {p:.3f}")
    if not beyond_meta:
        why.append("no significant gain over site/season/batch metadata")
    out.update(level=1, label="signal present but explained by site/season/batch", reason="; ".join(why))
    return out


def holm(pvals):
    """Holm step-down adjusted p-values for a dict name -> p (None entries ignored)."""
    items = [(k, p) for k, p in pvals.items() if p is not None]
    m = len(items)
    order = sorted(items, key=lambda kv: kv[1])
    adj, running = {}, 0.0
    for i, (k, p) in enumerate(order):
        running = max(running, (m - i) * p)
        adj[k] = min(1.0, running)
    return adj


# --------------------------------------------------------------------------- main
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--workbook", default=os.path.join(HERE, "AI_Data_Training.xlsx"))
    ap.add_argument("--audio-dir", default=None)
    ap.add_argument("--targets", nargs="*", default=None)
    ap.add_argument("--families", nargs="*", default=None)
    ap.add_argument("--n-repeats", type=int, default=N_REPEATS)
    ap.add_argument("--n-boot", type=int, default=N_BOOT)
    ap.add_argument("--n-perm", type=int, default=0, help="permutations per target (0 = skip)")
    ap.add_argument("--n-jobs", type=int, default=4)
    ap.add_argument("--out", default=RES_DIR)
    ap.add_argument("--exclude", nargs="*", default=(), help="feature tables to ignore (e.g. clap gem)")
    a = ap.parse_args()
    os.makedirs(os.path.join(a.out, "oof"), exist_ok=True)

    df = bd.load_workbook(a.workbook, audio_dir=a.audio_dir)
    df = df[df["has_audio"]].reset_index(drop=True) if a.audio_dir else df
    df, fam, names, cov = load_features(df, exclude=a.exclude)
    print(f"{len(df)} recordings with every feature table; families: "
          + ", ".join(f"{k}[{v.shape[1]}]" for k, v in fam.items()), flush=True)
    targets = bd.build_targets(df)
    if a.targets:
        targets = [t for t in targets if t["name"] in a.targets]
    families = a.families or [f for f in FAMILY_MODEL if f in fam]

    from joblib import Parallel, delayed
    print(f"evaluating {len(targets)} targets x {len(families)} families ...", flush=True)
    outs = Parallel(n_jobs=a.n_jobs, backend="loky")(
        delayed(evaluate_target)(t, fam, df, families, a.n_repeats, a.n_boot) for t in targets)
    results = {}
    for (res, oof), t in zip(outs, targets):
        results[t["name"]] = res
        if oof is not None:
            np.savez_compressed(os.path.join(a.out, "oof", f"{t['name']}.npz"), **oof)

    perm = {}
    if a.n_perm > 0:
        print(f"permutation tests ({a.n_perm} per target, month-aligned colony permutation) ...", flush=True)
        pf = [f for f in families if f in PERM_FAMILIES]
        pouts = Parallel(n_jobs=a.n_jobs, backend="loky")(
            delayed(permutation_target)(t, fam, df, pf, a.n_perm) for t in targets)
        perm = {t["name"]: p for t, p in zip(targets, pouts) if p}
        json.dump(perm, open(os.path.join(a.out, "permutation.json"), "w"), indent=1)
    for name, res in results.items():
        res["verdict"] = verdict(res, perm.get(name))
        res["permutation"] = perm.get(name)
    # Holm adjustment across the confirmatory (non-probe) targets, primary family
    pv = {n: (r["verdict"].get("p_perm")) for n, r in results.items() if not n.startswith("probe_")}
    adj = holm(pv)
    for n, r in results.items():
        if n in adj:
            r["verdict"]["p_perm_holm"] = adj[n]
            r["verdict"]["survives_holm"] = bool(adj[n] < 0.05)
    meta = dict(n_recordings=int(len(df)), n_colonies=int(df["colony"].nunique()),
                families={k: int(v.shape[1]) for k, v in fam.items()}, family_model=FAMILY_MODEL,
                feature_coverage=cov, n_splits=N_SPLITS, n_repeats=a.n_repeats, n_boot=a.n_boot, n_perm=a.n_perm,
                pca_k=PCA_K, lin_C=LIN_C, ridge_alpha=RIDGE_ALPHA, min_pos_colonies=MIN_POS_COLONIES,
                metadata_columns=names.get("meta"))
    json.dump(dict(meta=meta, results=results), open(os.path.join(a.out, "results.json"), "w"), indent=1, default=float)

    rows = []
    for name, res in results.items():
        for f, m in res.get("families", {}).items():
            rows.append(dict(target=name, kind=res["kind"], family=f, n=res["n"],
                             **{k: v for k, v in m.items() if k != "ci"}, ci_lo=m["ci"][0], ci_hi=m["ci"][1],
                             p_perm=((res.get("permutation") or {}).get(f) or {}).get("p_value"),
                             verdict=res["verdict"]["label"]))
    pd.DataFrame(rows).to_csv(os.path.join(a.out, "results_long.csv"), index=False)
    print("\nVERDICTS")
    for name, res in results.items():
        v = res["verdict"]
        val = f"{v['value']:.3f} [{v['ci'][0]:.3f}, {v['ci'][1]:.3f}]" if v.get("value") is not None else ""
        pp = f" p={v['p_perm']:.3f}" if v.get("p_perm") is not None else ""
        print(f"  {name:34s} {v['label']:55s} {val}{pp}")
    print("saved", a.out)


if __name__ == "__main__":
    main()