#!/usr/bin/env python3
"""
Multimodal theme classification - 5-fold cross-validation over every combination of
{lyrics, vocal, instrument, mix}. All settings (data, lyrics variant, grid, training) come
from classification/config.yaml.

Input: the single combined file built by the embedding stage
  data/multimodal/metadata.csv   songs, labels, CV 'fold' and 'lyrics_group'
  data/multimodal/embeddings.npz vocal / instrument / mix (n, 9984 = 13 layers x 768) and the
                                 lyrics variants: lyrics_phasea (n, 768),
                                 lyrics_phaseb_cv (5, n, 768: [k] trained without fold k)

Evaluation: for fold k the test songs are the fold's songs; an inner validation split of the
remaining songs (grouped by identical lyrics) is used for early stopping and to choose the
configuration; the test songs are only scored. The reported number per combination is the test
macro-F1 of the configuration chosen on validation, averaged over the 5 folds.

Fusion methods:
  mlp    per-modality projection (audio: learned softmax weights over its 13 hidden layers, as in
         SUPERB / MERT probing), concatenate, MLP head
  gmu    Gated Multimodal Unit (Arevalo et al., 2017): h_i = tanh(W_i x_i),
         z_i = sigmoid(W_zi [x_1..x_N]), h = sum_i z_i * h_i - per-song, per-feature gates
  xattn  gated multi-head cross-attention: the lyrics vector queries one token per hidden layer
         of every audio model (13 x n_audio tokens); a tanh gate initialised at 0 adds the audio
         context (cross-modal attention as in MulT, Tsai et al., 2019). Without lyrics a learned
         query token is used.
  late   average of the class probabilities of the best single-modality models
Training options (grid): modality dropout (ModDrop, Neverova et al., 2016), auxiliary
per-modality losses (motivated by Gradient Blending, Wang et al., 2020 / OGM-GE, Peng et al.,
2022; fixed weight here), soft labels (final_* theme scores instead of the arg-max label).

Outputs (output_dir from the config):
  runs.csv              every (fold, combo, configuration): val / test macro-F1
  summary.csv           per combo: mean/std test macro-F1 of the chosen model, per method too
  effects.csv           paired effect of each option / method (combo x fold)
  gain_over_lyrics.csv  per combo with lyrics + audio: test macro-F1 minus lyrics-only, per fold
  oof_predictions.npz   out-of-fold class probabilities of the chosen model per combo, plus
                        per-song GMU gate means of the all-modality model (used by the analysis)
  cv_comparison.png

Usage:
  python classification/classify.py                                    # config.yaml as is
  python classification/classify.py --set data.lyrics=lyrics_phasea \\
         --set output_dir=results/classification/phasea
  python classification/classify.py --set folds=[0] --set combos=[lyrics,lyrics+vocal] \\
         --set training.max_epochs=2 --set output_dir=/tmp/smoke                # smoke test
  CLASSIFY_SET="data.lyrics=lyrics_phasea;output_dir=results/classification/phasea" \\
      python classification/classify.py      # same overrides from the environment (Kaggle)
  python classification/classify.py --summarize-only                   # re-summarize runs.csv
"""

import argparse
import copy
import os
import time
from itertools import combinations, product

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
import yaml
from sklearn.metrics import accuracy_score, f1_score
from sklearn.model_selection import StratifiedGroupKFold

THEMES = ["romance", "sadness", "celebration", "motivational"]
AUDIO = {"vocal", "instrument", "mix"}
N_LAYERS, LAYER_DIM = 13, 768                 # audio vectors = 13 hidden layers x 768

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


# ═════════════════════════════ config ═════════════════════════════

def load_config(path, overrides):
    with open(path) as f:
        cfg = yaml.safe_load(f)
    for item in overrides:
        key, value = item.split("=", 1)
        node = cfg
        *parents, leaf = key.split(".")
        for p in parents:
            node = node.setdefault(p, {})
        node[leaf] = yaml.safe_load(value)
    return cfg


def resolve(path):
    """The path itself, or - on Kaggle - the same file name inside an attached dataset."""
    if os.path.exists(path):
        return path
    for root, _, files in os.walk("/kaggle/input"):
        if os.path.basename(path) in files:
            return os.path.join(root, os.path.basename(path))
    raise SystemExit(f"{path} not found (also searched /kaggle/input)")


# ═════════════════════════════ data ═════════════════════════════

def load_data(cfg):
    meta = pd.read_csv(resolve(cfg["data"]["metadata"]), low_memory=False)
    npz = np.load(resolve(cfg["data"]["embeddings"]))
    audio = {m: npz[m][meta[f"{m}_embedding_idx"].to_numpy()].astype(np.float32)
             for m in cfg["modalities"] if m in AUDIO}
    lyrics = npz[cfg["data"]["lyrics"]].astype(np.float32)
    rows = meta["lyrics_embedding_idx"].to_numpy()
    lyrics = lyrics[:, rows] if lyrics.ndim == 3 else lyrics[rows]
    folds = meta["fold"].to_numpy()
    if lyrics.ndim == 3:
        assert lyrics.shape[0] == len(set(folds)), "one lyrics array per fold expected"

    y = meta["dominant_theme"].map({t: i for i, t in enumerate(THEMES)}).to_numpy()
    soft = meta[[f"final_{t}" for t in THEMES]].to_numpy(dtype=np.float32)
    zero = soft.sum(1) <= 0
    soft[zero] = np.eye(len(THEMES), dtype=np.float32)[y[zero]]
    soft /= soft.sum(1, keepdims=True)
    return meta, audio, lyrics, folds, meta["lyrics_group"].to_numpy(), y, soft


def inner_split(idx, y, groups, cfg):
    """Hold out 1/inner_val_splits of the non-test songs (grouped, stratified) for validation."""
    sgkf = StratifiedGroupKFold(n_splits=cfg["inner_val_splits"], shuffle=True,
                                random_state=cfg["seed"])
    tr, va = next(sgkf.split(idx, y[idx], groups[idx]))
    return idx[tr], idx[va]


def fold_tensors(k, audio, lyrics, split):
    """Standardize every modality on this fold's TRAIN rows; tensors for train/val/test."""
    sources = {**audio, "lyrics": lyrics[k] if lyrics.ndim == 3 else lyrics}
    out = {}
    for m, X in sources.items():
        mu = X[split["train"]].mean(0)
        sd = X[split["train"]].std(0)
        sd[sd < 1e-8] = 1.0
        out[m] = {s: torch.tensor((X[ix] - mu) / sd, dtype=torch.float32, device=device)
                  for s, ix in split.items()}
    return out


# ═════════════════════════════ models ═════════════════════════════

class LayerWeightedSum(nn.Module):
    """Learned softmax weights over the 13 hidden layers -> one 768-dim vector."""
    def __init__(self):
        super().__init__()
        self.w = nn.Parameter(torch.zeros(N_LAYERS))

    def forward(self, x):
        x = x.view(-1, N_LAYERS, LAYER_DIM)
        return (torch.softmax(self.w, 0)[None, :, None] * x).sum(1)


def modality_encoder(in_dim, d_model, dropout, layered):
    head = [LayerWeightedSum()] if layered else []
    return nn.Sequential(*head, nn.Linear(LAYER_DIM if layered else in_dim, d_model),
                         nn.LayerNorm(d_model), nn.GELU(), nn.Dropout(dropout))


def keep_mask(batch, k, p, dev):
    """Modality dropout: drop each modality with prob p, but never all of them for a song."""
    keep = (torch.rand(batch, k, device=dev) >= p).float()
    none = keep.sum(1) == 0
    if none.any():
        keep[none, torch.randint(k, (int(none.sum()),), device=dev)] = 1.0
    return keep


class UnimodalAux(nn.Module):
    """Auxiliary per-modality classifiers on the encoded features. Their losses keep every
    modality's encoder learning even when lyrics dominate the fused loss."""
    def __init__(self, k, d_model):
        super().__init__()
        self.heads = nn.ModuleList(nn.Linear(d_model, len(THEMES)) for _ in range(k))
        self.last = None

    def __call__(self, feats):
        self.last = [h(f) for h, f in zip(self.heads, feats)]


class FusionMLP(nn.Module):
    def __init__(self, mods, in_dims, d_model, dropout, moddrop=0.0, aux=False):
        super().__init__()
        self.moddrop = moddrop
        self.enc = nn.ModuleList(modality_encoder(d, d_model, dropout, m in AUDIO)
                                 for m, d in zip(mods, in_dims))
        k = len(mods)
        self.aux = UnimodalAux(k, d_model) if aux and k > 1 else None
        self.head = nn.Sequential(nn.Linear(k * d_model, d_model // 2), nn.LayerNorm(d_model // 2),
                                  nn.GELU(), nn.Dropout(dropout), nn.Linear(d_model // 2, len(THEMES)))

    def forward(self, xs):
        h = [e(x) for e, x in zip(self.enc, xs)]
        if self.aux is not None and self.training:
            self.aux(h)
        if self.training and self.moddrop > 0 and len(h) > 1:
            keep = keep_mask(len(h[0]), len(h), self.moddrop, h[0].device)
            h = [hi * keep[:, i:i + 1] for i, hi in enumerate(h)]
        return self.head(torch.cat(h, dim=1))


class FusionGMU(nn.Module):
    """Gated Multimodal Unit, N-modality form. The gates z_i are per song and per feature, so
    their mean is a per-song measure of how much the model relies on each modality."""
    def __init__(self, mods, in_dims, d_model, dropout, moddrop=0.0, aux=False):
        super().__init__()
        self.moddrop = moddrop
        k = len(mods)
        self.enc = nn.ModuleList(modality_encoder(d, d_model, dropout, m in AUDIO)
                                 for m, d in zip(mods, in_dims))
        self.hid = nn.ModuleList(nn.Linear(d_model, d_model) for _ in range(k))
        self.gate = nn.Linear(k * d_model, k * d_model)
        self.aux = UnimodalAux(k, d_model) if aux and k > 1 else None
        self.head = nn.Sequential(nn.LayerNorm(d_model), nn.Dropout(dropout),
                                  nn.Linear(d_model, d_model // 2), nn.GELU(), nn.Dropout(dropout),
                                  nn.Linear(d_model // 2, len(THEMES)))

    def fuse(self, xs):
        x = [e(xi) for e, xi in zip(self.enc, xs)]
        if self.aux is not None and self.training:
            self.aux(x)
        k, b = len(x), len(x[0])
        if self.training and self.moddrop > 0 and k > 1:
            keep = keep_mask(b, k, self.moddrop, x[0].device)
            x = [xi * keep[:, i:i + 1] for i, xi in enumerate(x)]
        h = torch.stack([torch.tanh(w(xi)) for w, xi in zip(self.hid, x)], dim=1)     # (b, k, d)
        z = torch.sigmoid(self.gate(torch.cat(x, dim=1))).view(b, k, -1)             # (b, k, d)
        return (z * h).sum(1), z

    def forward(self, xs):
        return self.head(self.fuse(xs)[0])

    @torch.no_grad()
    def gate_means(self, xs, bs=512):
        """Per-song mean gate value per modality, (n, k)."""
        self.eval()
        return torch.cat([self.fuse([x[i:i + bs] for x in xs])[1].mean(2)
                          for i in range(0, len(xs[0]), bs)]).cpu().numpy()


class GatedCrossAttention(nn.Module):
    def __init__(self, mods, in_dims, d_model, dropout, nhead, moddrop=0.0):
        super().__init__()
        self.moddrop = moddrop
        self.text_i = [i for i, m in enumerate(mods) if m not in AUDIO]
        self.audio_i = [i for i, m in enumerate(mods) if m in AUDIO]
        assert self.audio_i, "cross-attention needs at least one audio modality"
        self.q_enc = (modality_encoder(in_dims[self.text_i[0]], d_model, dropout, False)
                      if self.text_i else None)
        self.q_null = nn.Parameter(torch.zeros(1, d_model))      # query when lyrics absent/dropped
        self.a_proj = nn.ModuleList(nn.Sequential(nn.Linear(LAYER_DIM, d_model), nn.LayerNorm(d_model))
                                    for _ in self.audio_i)
        self.layer_emb = nn.Parameter(torch.zeros(1, N_LAYERS, d_model))
        self.mod_emb = nn.Parameter(torch.zeros(len(self.audio_i), 1, d_model))
        for p in (self.q_null, self.layer_emb, self.mod_emb):
            nn.init.trunc_normal_(p, std=0.02)
        self.attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
        self.gate = nn.Parameter(torch.zeros(1)) if self.text_i else None
        self.ffn = nn.Sequential(nn.LayerNorm(d_model), nn.Linear(d_model, 2 * d_model), nn.GELU(),
                                 nn.Dropout(dropout), nn.Linear(2 * d_model, d_model))
        self.head = nn.Sequential(nn.LayerNorm(d_model), nn.Linear(d_model, d_model // 2), nn.GELU(),
                                  nn.Dropout(dropout), nn.Linear(d_model // 2, len(THEMES)))

    def forward(self, xs):
        b = len(xs[0])
        kv = torch.cat([proj(xs[i].view(b, N_LAYERS, LAYER_DIM)) + self.layer_emb + self.mod_emb[j]
                        for j, (i, proj) in enumerate(zip(self.audio_i, self.a_proj))], dim=1)

        n_mod = len(self.text_i) + len(self.audio_i)          # text column first, then audio
        keep = (keep_mask(b, n_mod, self.moddrop, kv.device)
                if self.training and self.moddrop > 0 and n_mod > 1
                else torch.ones(b, n_mod, device=kv.device))
        if self.q_enc is not None:
            kt = keep[:, :1]
            q = self.q_enc(xs[self.text_i[0]]) * kt + self.q_null * (1 - kt)
            akeep = keep[:, 1:]
        else:
            q = self.q_null.expand(b, -1)
            akeep = keep

        pad = (akeep == 0).repeat_interleave(N_LAYERS, dim=1)     # True = ignore this token
        no_audio = pad.all(1)
        if no_audio.any():                                     # avoid an all-masked softmax (NaN)
            pad = pad.clone()
            pad[no_audio, 0] = False
        ctx = self.attn(q[:, None], kv, kv, key_padding_mask=pad, need_weights=False)[0][:, 0]
        ctx = ctx * (~no_audio).float()[:, None]

        h = q + (torch.tanh(self.gate) * ctx if self.gate is not None else ctx)
        return self.head(h + self.ffn(h))


def build_model(run, mods, in_dims, mcfg):
    if run["method"] == "xattn":
        return GatedCrossAttention(mods, in_dims, mcfg["d_model"], mcfg["dropout"],
                                   mcfg["attention_heads"], moddrop=run["moddrop"]).to(device)
    cls = FusionGMU if run["method"] == "gmu" else FusionMLP
    return cls(mods, in_dims, mcfg["d_model"], mcfg["dropout"], moddrop=run["moddrop"],
               aux=run["aux"]).to(device)


# ═════════════════════════════ training ═════════════════════════════

@torch.no_grad()
def predict_proba(model, xs, bs=512):
    model.eval()
    return torch.cat([F.softmax(model([x[i:i + bs] for x in xs]), 1)
                      for i in range(0, len(xs[0]), bs)]).cpu().numpy()


def train_eval(run, mods, T, y, soft, split, cfg, want_gates=False):
    tcfg = cfg["training"]
    torch.manual_seed(cfg["seed"])
    np.random.seed(cfg["seed"])
    X = {s: [T[m][s] for m in mods] for s in split}
    ytr = torch.tensor(y[split["train"]], device=device)
    str_ = torch.tensor(soft[split["train"]], device=device)
    model = build_model(run, mods, [x.shape[1] for x in X["train"]], cfg["model"])
    opt = torch.optim.AdamW(model.parameters(), lr=run["lr"], weight_decay=tcfg["weight_decay"])
    sched = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, mode="max", factor=0.5,
                                                       patience=tcfg["lr_scheduler_patience"])

    def objective(logits, b):
        if run["soft"]:
            return -(str_[b] * F.log_softmax(logits, 1)).sum(1).mean()
        return F.cross_entropy(logits, ytr[b])

    best, best_state, bad, t0, n = -1.0, None, 0, time.time(), len(ytr)
    epoch = 0
    for epoch in range(1, tcfg["max_epochs"] + 1):
        model.train()
        for b in torch.randperm(n, device=device).split(tcfg["batch_size"]):
            loss = objective(model([x[b] for x in X["train"]]), b)
            aux = getattr(model, "aux", None)
            if aux is not None:
                loss = loss + cfg["model"]["aux_loss_weight"] * torch.stack(
                    [objective(a, b) for a in aux.last]).mean()
            opt.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
        val_f1 = f1_score(y[split["val"]], predict_proba(model, X["val"]).argmax(1), average="macro")
        sched.step(val_f1)
        if val_f1 > best + tcfg["min_delta"]:
            best, bad, best_state = val_f1, 0, copy.deepcopy(model.state_dict())
        else:
            bad += 1
        if epoch >= min(tcfg["min_epochs"], tcfg["max_epochs"]) and bad >= tcfg["patience"]:
            break
    model.load_state_dict(best_state)
    pv, pt = predict_proba(model, X["val"]), predict_proba(model, X["test"])
    gate = (float(torch.tanh(model.gate).item())
            if isinstance(model, GatedCrossAttention) and model.gate is not None else np.nan)
    gates = model.gate_means(X["test"]) if want_gates and isinstance(model, FusionGMU) else None
    return {"val_f1": best, "pv": pv, "pt": pt, "epochs": epoch, "seconds": time.time() - t0,
            "gate": gate, "gmu_gates": gates, "model": model}


def runs_for(mods, grid):
    """The grid for one combination. Single modality: no modality dropout / aux losses, and no
    cross-attention for lyrics alone (it needs audio tokens to attend to)."""
    if len(mods) == 1:
        methods = [m for m in grid["single_modality_methods"]
                   if not (m == "xattn" and mods[0] not in AUDIO)]
    else:
        methods = grid["methods"]
    out = []
    for m in methods:
        moddrop = grid["modality_dropout"] if len(mods) > 1 else [0.0]
        aux = grid["aux_loss"] if len(mods) > 1 and m != "xattn" else [False]
        out += [{"method": m, "moddrop": d, "aux": a, "soft": s, "lr": lr}
                for d, a, s, lr in product(moddrop, aux, grid["soft_labels"], grid["lr"])]
    return out


def scores(y_true, proba):
    p = proba.argmax(1)
    return f1_score(y_true, p, average="macro"), accuracy_score(y_true, p)


def run_fold(k, combos, audio, lyrics, folds, groups, y, soft, cfg):
    """Returns the run rows, the out-of-fold test probabilities of the configuration chosen on
    validation for every combo, and per-song GMU gates for the all-modality combo."""
    test = np.flatnonzero(folds == k)
    tr, va = inner_split(np.flatnonzero(folds != k), y, groups, cfg)
    split = {"train": tr, "val": va, "test": test}
    print(f"\n{'=' * 70}\nFOLD {k}: train {len(tr)} | val {len(va)} | test {len(test)}\n{'=' * 70}",
          flush=True)
    T = fold_tensors(k, audio, lyrics, split)
    yv, yt = y[va], y[test]
    all_name = "+".join(cfg["modalities"])

    rows, probs, best_single, best_gmu = [], {}, {}, None
    for mods in combos:
        name = "+".join(mods)
        for run in runs_for(mods, cfg["grid"]):
            r = train_eval(run, mods, T, y, soft, split, cfg, want_gates=(name == all_name))
            tf1, tacc = scores(yt, r["pt"])
            rows.append({"fold": k, "combo": name, "n_mod": len(mods), **run,
                         "val_f1": r["val_f1"], "test_f1": tf1, "test_acc": tacc,
                         "epochs": r["epochs"], "seconds": r["seconds"], "gate": r["gate"]})
            probs[len(rows) - 1] = r["pt"]
            if len(mods) == 1 and r["val_f1"] > best_single.get(name, {"val_f1": -1})["val_f1"]:
                best_single[name] = r
            if r["gmu_gates"] is not None and (best_gmu is None or r["val_f1"] > best_gmu["val_f1"]):
                best_gmu = r
        top = max((x for x in rows if x["combo"] == name), key=lambda x: x["val_f1"])
        print(f"  {name:32s} chosen: {top['method']:5s} moddrop={top['moddrop']} "
              f"aux={top['aux']!s:5s} soft={top['soft']!s:5s} lr={top['lr']:.0e} | "
              f"val {top['val_f1']:.4f} test {top['test_f1']:.4f}", flush=True)

    if cfg.get("late_fusion", True):
        for mods in combos:
            if len(mods) > 1 and all(m in best_single for m in mods):
                pv = np.mean([best_single[m]["pv"] for m in mods], axis=0)
                pt = np.mean([best_single[m]["pt"] for m in mods], axis=0)
                vf1, _ = scores(yv, pv)
                tf1, tacc = scores(yt, pt)
                rows.append({"fold": k, "combo": "+".join(mods), "n_mod": len(mods),
                             "method": "late", "moddrop": np.nan, "aux": np.nan, "soft": np.nan,
                             "lr": np.nan, "val_f1": vf1, "test_f1": tf1, "test_acc": tacc,
                             "epochs": np.nan, "seconds": 0.0, "gate": np.nan})
                probs[len(rows) - 1] = pt

    oof = {}
    for name in dict.fromkeys(r["combo"] for r in rows):          # same rule as pick_by_val
        i = max((j for j, r in enumerate(rows) if r["combo"] == name), key=lambda j: rows[j]["val_f1"])
        oof[name] = probs[i]
    gates = best_gmu["gmu_gates"] if best_gmu is not None else None
    return rows, test, oof, gates


# ═════════════════════════════ summary ═════════════════════════════

def pick_by_val(df):
    """Per (combo, fold): the configuration with the best validation F1 - the chosen model."""
    return df.loc[df.groupby(["combo", "fold"])["val_f1"].idxmax()]


def paired(a, b, name):
    d = (a - b).dropna()
    return {"comparison": name, "mean_test_f1_gain": d.mean(), "std": d.std(ddof=0),
            "wins": int((d > 0).sum()), "n": len(d)}


def summarize(out_dir):
    runs = pd.read_csv(os.path.join(out_dir, "runs.csv"))
    n_folds = runs["fold"].nunique()
    chosen = pick_by_val(runs)

    def agg(sel, label):
        g = sel.groupby("combo")["test_f1"]
        return pd.DataFrame({f"{label}_mean": g.mean(), f"{label}_std": g.std(ddof=0)})

    parts = [agg(chosen, "chosen")]
    for m in ["mlp", "gmu", "xattn", "late"]:
        sub = runs[runs["method"] == m]
        if len(sub):
            parts.append(agg(pick_by_val(sub), m))
    summ = pd.concat(parts, axis=1)
    summ["chosen_method"] = chosen.groupby("combo")["method"].agg(lambda s: s.value_counts().index[0])
    summ["n_mod"] = runs.groupby("combo")["n_mod"].first()
    summ = summ.sort_values("chosen_mean", ascending=False)
    summ.to_csv(os.path.join(out_dir, "summary.csv"))

    # paired effects of each option / method (best-by-val with vs without), multi-modality combos
    multi = runs[(runs["n_mod"] > 1) & runs["method"].isin(["mlp", "gmu", "xattn"])]
    by = ["combo", "fold"]
    eff = []
    for opt, on, off in [("moddrop", 0.3, 0.0), ("soft", True, False), ("aux", True, False)]:
        sub = multi[multi["method"].isin(["mlp", "gmu"])] if opt == "aux" else multi
        col = sub[opt].astype(str)
        eff.append(paired(pick_by_val(sub[col == str(on)]).set_index(by)["test_f1"],
                          pick_by_val(sub[col == str(off)]).set_index(by)["test_f1"], opt))
    ml = pick_by_val(multi[multi["method"] == "mlp"]).set_index(by)["test_f1"]
    for m in ["gmu", "xattn"]:
        eff.append(paired(pick_by_val(multi[multi["method"] == m]).set_index(by)["test_f1"], ml,
                          f"{m}_vs_mlp"))
    late = runs[(runs["n_mod"] > 1) & (runs["method"] == "late")]
    if len(late):
        eff.append(paired(pick_by_val(late).set_index(by)["test_f1"], ml, "late_vs_mlp"))
    eff = pd.DataFrame(eff)
    eff.to_csv(os.path.join(out_dir, "effects.csv"), index=False)

    # does adding audio to lyrics help? chosen model of each lyrics+audio combo vs lyrics alone
    ch = chosen.set_index(by)["test_f1"]
    gain = []
    if "lyrics" in set(runs["combo"]):
        base = ch.xs("lyrics", level="combo")
        for combo in summ.index:
            parts_ = combo.split("+")
            if "lyrics" in parts_ and len(parts_) > 1:
                g = paired(ch.xs(combo, level="combo"), base, combo)
                g.update({f"fold_{f}": v for f, v in (ch.xs(combo, level="combo") - base).items()})
                gain.append(g)
    gain = pd.DataFrame(gain)
    if len(gain):
        gain.to_csv(os.path.join(out_dir, "gain_over_lyrics.csv"), index=False)

    pd.set_option("display.width", 220)
    print(f"\n=== {n_folds}-fold CV: test macro-F1 of the configuration chosen on validation ===")
    print(summ.round(4).to_string())
    print("\n=== Paired effect of each option / method (combo x fold, multi-modality combos) ===")
    print(eff.round(4).to_string(index=False))
    if len(gain):
        print("\n=== Lyrics + audio vs lyrics alone (per fold) ===")
        print(gain.round(4).to_string(index=False))

    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    ink, muted, grid_c, bar_c, surface = "#0b0b0b", "#898781", "#e1e0d9", "#2a78d6", "#fcfcfb"
    d = summ.iloc[::-1]
    labels = [c + ("  (old, leaky)" if c == "lyrics_leaky" else "") for c in d.index]
    fig, ax = plt.subplots(figsize=(8.5, 0.38 * len(d) + 1.4), facecolor=surface)
    ax.set_facecolor(surface)
    ax.barh(range(len(d)), d["chosen_mean"], height=0.6, color=bar_c,
            xerr=d["chosen_std"], error_kw={"ecolor": ink, "elinewidth": 1, "capsize": 2})
    ax.set_yticks(range(len(d)), labels, fontsize=9, color=ink)
    for i, (v, s) in enumerate(zip(d["chosen_mean"], d["chosen_std"])):
        ax.text(v + s + 0.01, i, f"{v:.3f}", va="center", fontsize=8, color=muted)
    ax.set_xlim(0, 1.0)
    ax.axvline(1 / len(THEMES), color=muted, linewidth=1, linestyle=":")
    ax.text(1 / len(THEMES), len(d) - 0.4, " chance", color=muted, fontsize=8, va="bottom")
    ax.set_xlabel(f"Test macro-F1 ({n_folds}-fold CV mean, bars = ±1 std across folds)", color=muted)
    ax.set_title("Cross-validated test macro-F1 by modality combination", loc="left",
                 color=ink, fontsize=11)
    ax.grid(axis="x", color=grid_c, linewidth=0.8)
    ax.set_axisbelow(True)
    for s in ("top", "right", "left"):
        ax.spines[s].set_visible(False)
    ax.spines["bottom"].set_color(grid_c)
    ax.tick_params(colors=muted, length=0)
    plt.tight_layout()
    plt.savefig(os.path.join(out_dir, "cv_comparison.png"), dpi=150, facecolor=surface)
    plt.close(fig)
    print(f"\nSaved summary.csv, effects.csv, gain_over_lyrics.csv, cv_comparison.png -> {out_dir}")


# ═════════════════════════════ main ═════════════════════════════

def main():
    ap = argparse.ArgumentParser(description="5-fold CV multimodal theme classification")
    ap.add_argument("--config", default="classification/config.yaml")
    ap.add_argument("--set", action="append", default=[], metavar="KEY=VALUE",
                    help="override a config value, e.g. --set data.lyrics=lyrics_phasea")
    ap.add_argument("--summarize-only", action="store_true")
    args = ap.parse_args()
    # Kaggle script kernels take no arguments: CLASSIFY_SET="key=value;key=value" (push_to_kaggle.sh SET_ENV)
    env_set = [kv for kv in os.environ.get("CLASSIFY_SET", "").split(";") if kv]
    cfg = load_config(args.config, env_set + args.set)
    out_dir = cfg["output_dir"]
    if args.summarize_only:
        summarize(out_dir)
        return

    meta, audio, lyrics, folds, groups, y, soft = load_data(cfg)
    gpu = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "no GPU"
    print(f"Device: {device} ({gpu}) | torch {torch.__version__} | songs {len(y)} | "
          f"lyrics: {cfg['data']['lyrics']} | folds {np.bincount(folds).tolist()}", flush=True)

    mods = cfg["modalities"]
    if cfg["combos"] == "all":
        combos = [list(c) for k in range(1, len(mods) + 1) for c in combinations(mods, k)]
    else:
        combos = [[m for m in mods if m in c.split("+")] for c in cfg["combos"]]
    combos.sort(key=len)                     # single modalities first: late fusion needs them
    fold_ids = sorted(set(folds)) if cfg["folds"] == "all" else list(cfg["folds"])

    os.makedirs(out_dir, exist_ok=True)
    with open(os.path.join(out_dir, "config_used.yaml"), "w") as f:
        yaml.safe_dump(cfg, f, sort_keys=False)
    rows, oof = [], {}
    gates = np.full((len(y), len(mods)), np.nan, dtype=np.float32)
    for k in fold_ids:
        fold_rows, test, fold_oof, fold_gates = run_fold(k, combos, audio, lyrics, folds, groups,
                                                         y, soft, cfg)
        rows += fold_rows
        for name, p in fold_oof.items():
            oof.setdefault(name, np.full((len(y), len(THEMES)), np.nan, dtype=np.float32))[test] = p
        if fold_gates is not None:
            gates[test] = fold_gates
        pd.DataFrame(rows).to_csv(os.path.join(out_dir, "runs.csv"), index=False)
        np.savez(os.path.join(out_dir, "oof_predictions.npz"), video_id=meta["video_id"].to_numpy(str),
                 y=y, fold=folds, gmu_gates_all=gates, **oof)
    summarize(out_dir)


if __name__ == "__main__":
    main()