"""Final protocol: every head is trained on REAL pairs only, with an explicit patient-level train / val / test split per
real dataset; virtual modalities (STPath RNA, GigaTIME mIF, GPT synthetic captions) are used ONLY at test time.

Real datasets and splits (5 folds each; fold k = train / val (early stopping only) / test; patients never cross splits)
  ORION  H&E <-> mIF  : splits/patient_5fold/fold<k>.csv (5-fold by patient, val = 3 training patients)
                        + splits/orion_official/fold0.csv (the user's ORION train/val/test csvs)
  HEST   H&E <-> RNA  : GroupKFold(5) over study|patient groups, val = 10% of training groups (= 23_bind_gene.py)
  TCGA   H&E <-> text : GroupKFold(5) over patients on ALL image-text pairs, val = 10% of training patients (= 20_pairing.py)
The HEST and TCGA split memberships are written to splits/final/{hest,tcga}_fold<k>.csv.

Heads (frozen UNI2 H&E = anchor space):
  mIF  runs/bench_fm/<split>_fold<k>_uni2_fm-virtues[_shared]_.../heads.pt       (17_bench_fm.py --save_heads)
  RNA  runs/gene/uni2_mlp_fold<k>/head.pt                                         (23_bind_gene.py)
  text runs/text_clean/uni2_biomedclip_fold<k>/head.pt trained here (text_head recipe of 20_pairing.py) on CLEAN real
       texts only: junk texts (answer letters, "error", < 5 words; 40_clean_text.py) are removed from train/val/test.
       The 256 images left without real text keep their split; their few-shot GPT captions (27_ --sets tcga_notext)
       are evaluated only (pooled over the test folds, each image through its own test-fold head).
       Previous heads trained with the junk texts: runs/text_all (results in results/before_text_clean/).
Evaluation is on the TEST split of fold k with the fold-k heads (a head from another dataset never saw that dataset).
  part 1 "real"    : both sides measured           (H&E<->mIF, H&E<->RNA, H&E<->text)
  part 2 "virtual" : >= 1 side virtual, same test items; "circular" = both sides derived from the same H&E
Metrics: random 100-item galleries (50 draws), both directions averaged: tie-aware R@1/R@5/R@10, MRR, and one-to-one
Hungarian pairing accuracy (chance R@1 = pair acc = 0.01, R@5 = 0.05, R@10 = 0.10).
mIF everywhere = the 10-channel panel shared by ORION and GigaTIME (30_virtues_shared.py); the full 17-channel ORION
panel appears only as a reference row. "[qm]" = GigaTIME pixels quantile-mapped per channel onto the real fold-k
training-tile intensity distribution before the identical preprocessing (30_ virtual_qm).
Outputs results/final_long.csv, results/final_summary.csv; split sizes -> results/final_splits.csv
"""
import os

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from scipy.optimize import linear_sum_assignment
from sklearn.model_selection import GroupKFold

from common import CACHE, ROOT, RUNS, SPLITS, load_all_tiles

OUT = os.path.join(ROOT, "results")
H = os.path.join(CACHE, "hest")
G100, DRAWS, K = 100, 50, 5
dev = "cuda" if torch.cuda.is_available() else "cpu"
os.makedirs(os.path.join(SPLITS, "final"), exist_ok=True)


def mlp(din, dout=1536, hid=512):  # same head as 12_/17_/20_/23_
    return nn.Sequential(nn.Linear(din, hid), nn.GELU(), nn.Dropout(0.1), nn.Linear(hid, hid), nn.GELU(),
                         nn.LayerNorm(hid), nn.Linear(hid, dout))


def load_mlp(sd):
    h = mlp(sd["0.weight"].shape[1], sd[f"{len(mlp(1)) - 1}.weight"].shape[0]); h.load_state_dict(sd); return h.eval()


def nz(x):
    return F.normalize(torch.as_tensor(np.asarray(x), dtype=torch.float32), dim=-1)


def g100(a, b, seed=0):
    rng = np.random.default_rng(seed)
    out = {m: [] for m in ["R@1", "R@5", "R@10", "MRR", "pair_acc"]}
    for _ in range(DRAWS):
        s = torch.from_numpy(rng.choice(len(a), min(G100, len(a)), replace=False))
        sim = (a[s] @ b[s].T).numpy()
        d = np.diag(sim)
        for S in (sim, sim.T):
            rank = (S > d[:, None]).sum(1)  # tie-aware: 0 = correct item ranked first
            for k_ in (1, 5, 10):
                out[f"R@{k_}"].append((rank < k_).mean() / 2)
            out["MRR"].append((1 / (rank + 1)).mean() / 2)
        ri, ci = linear_sum_assignment(-sim)
        out["pair_acc"].append((ci[np.argsort(ri)] == np.arange(len(s))).mean())
    return {m: float(np.sum(v) / DRAWS) for m, v in out.items()}


rows, split_rows = [], []


def add(dataset, split, fold, part, pair, za, zb, kind, circular):
    rows.append({"dataset": dataset, "split": split, "fold": fold, "part": part, "pair": pair, "kind": kind,
                 "circular": circular, "n_test": len(za), **g100(za, zb)})


def log_split(dataset, split, fold, s, unit):
    for x in ["train", "val", "test"]:
        split_rows.append({"dataset": dataset, "split": split, "fold": fold, "role": x,
                           "n_items": int((s.split == x).sum()), f"n_{unit}": s[unit][s.split == x].nunique(),
                           "unit": unit})
    g = {x: set(s[unit][s.split == x]) for x in ["train", "val", "test"]}
    assert not (g["train"] & g["val"] or g["train"] & g["test"] or g["val"] & g["test"]), (dataset, fold, "leak")


# ================= virtual RNA: STPath counts -> log1p CP10k on the 2000 HEST HVGs (input of the HEST-trained head)
hvg = pd.read_csv(os.path.join(H, "hvg_genes.csv")).gene.values


def virtual_rna(st):
    g = pd.read_csv(os.path.join(CACHE, "stpath", f"{st}_genes.csv")).gene.values
    c = np.expm1(np.load(os.path.join(CACHE, "stpath", f"{st}_pred.npy")).astype(np.float32)).clip(0)
    c = np.log1p(c / c.sum(1, keepdims=True) * 1e4)
    return torch.from_numpy(c[:, pd.Series(range(len(g)), index=g).loc[hvg].values])


def mif_head(enc, k, split):
    d = torch.load(os.path.join(RUNS, "bench_fm", f"{split}_fold{k}_uni2_fm-{enc}_fm-geneformer_pathgenclip", "heads.pt"),
                   map_location="cpu")
    h = load_mlp(d["state"]["if"])
    return lambda x: h((torch.as_tensor(x, dtype=torch.float32) - d["if_mean"]) / d["if_std"])


rna_heads = [load_mlp(torch.load(os.path.join(RUNS, "gene", f"uni2_mlp_fold{k}", "head.pt"), map_location="cpu")) for k in range(K)]

# ================= TCGA all image-text pairs: split + text heads (real text only)
I = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv"))
T = pd.read_csv(os.path.join(CACHE, "alltxt_meta.csv"))
JUNK = pd.read_csv(os.path.join(CACHE, "alltxt_junk.csv")).junk.values
assert len(JUNK) == len(T)
IMG = torch.from_numpy(np.load(os.path.join(CACHE, "allimg_he_uni2.npy")).astype(np.float32))       # (n_img, 8, 1536)
TXT = torch.from_numpy(np.load(os.path.join(CACHE, "alltxt_text_biomedclip.npy"))[:, 0].astype(np.float32))
t_img = torch.as_tensor(T.img_idx.values)  # indexed by the ORIGINAL text row (kept after filtering)
T = T[~JUNK]
NOTXT = pd.read_csv(os.path.join(CACHE, "tcga_notext_meta.csv"))
assert not NOTXT.img_idx.isin(T.img_idx).any()
TX_N = torch.from_numpy(np.load(os.path.join(CACHE, "tcga_notext_text_biomedclip.npy"))[:, 0].astype(np.float32))


def one_text_per_image(t_ids, rng):
    sub = T.loc[t_ids]
    pick = sub.groupby("img_idx").apply(lambda g: g.index[rng.integers(len(g))], include_groups=False)
    return pick.index.values, pick.values  # img ids, text ids


def train_text_head(tr_t, va_t, epochs=60, bs=256, lr=3e-4, seed=0):
    """20_pairing.py text_head: frozen UNI2 image = anchor, MLP on frozen BiomedCLIP text, multi-positive InfoNCE"""
    torch.manual_seed(seed); rng = np.random.default_rng(seed)
    X, Tx, ti = IMG.to(dev), TXT.to(dev), t_img.to(dev)
    th = mlp(Tx.shape[1], X.shape[-1]).to(dev)
    logit = nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32, device=dev))
    opt = torch.optim.AdamW(list(th.parameters()) + [logit], lr=lr, weight_decay=0.05)
    sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=lr, total_steps=epochs * max(1, len(tr_t) // bs), pct_start=0.05)
    va_i, va_tt = one_text_per_image(va_t, np.random.default_rng(0))
    za_va = F.normalize(X[torch.as_tensor(va_i, device=dev), 0], dim=-1)
    best, state = -1, None
    for ep in range(epochs):
        th.train()
        perm = rng.permutation(tr_t)
        for s in range(0, len(perm) - bs + 1, bs):
            b = torch.as_tensor(perm[s:s + bs], device=dev)
            ii = ti[b]
            zi = F.normalize(X[ii, torch.randint(0, 8, (len(b),), device=dev)], dim=-1)
            lg = logit.exp().clamp(max=100) * zi @ F.normalize(th(Tx[b]), dim=-1).T
            pos = (ii[:, None] == ii[None, :]).float()
            tgt = pos / pos.sum(1, keepdim=True)
            loss = (torch.sum(-tgt * F.log_softmax(lg, 1), 1).mean() + torch.sum(-tgt * F.log_softmax(lg.T, 1), 1).mean()) / 2
            opt.zero_grad(set_to_none=True); loss.backward(); opt.step(); sched.step()
        th.eval()
        with torch.no_grad():
            zt = F.normalize(th(Tx[torch.as_tensor(va_tt, device=dev)]), dim=-1)
            sim = za_va @ zt.T
            v = float((((sim > sim.diag()[:, None]).sum(1) == 0).float().mean() + ((sim.T > sim.diag()[:, None]).sum(1) == 0).float().mean()) / 2)
        if v > best:
            best, state = v, {k_: x.detach().cpu().clone() for k_, x in th.state_dict().items()}
    return state, best


tcga_folds = list(GroupKFold(K).split(I, groups=I.patient))
text_heads, tcga_test, tcga_notext_test = [], [], []
for k, (trv, te) in enumerate(tcga_folds):
    pats = np.random.default_rng(k).permutation(sorted(I.patient.iloc[trv].unique()))
    va_p = set(pats[: max(1, len(pats) // 10)])
    sp = pd.Series("test", index=I.index)
    sp.iloc[trv] = np.where(I.patient.iloc[trv].isin(va_p), "val", "train")
    s = pd.DataFrame({"img_idx": I.index, "patient": I.patient, "split": sp.values})
    s.to_csv(os.path.join(SPLITS, "final", f"tcga_fold{k}.csv"), index=False)
    log_split("TCGA image-text", "patient 5-fold", k, s, "patient")
    tsplit = sp.loc[T.img_idx].values
    tr_t, va_t, te_t = (T.index[tsplit == x].values for x in ["train", "val", "test"])
    f = os.path.join(RUNS, "text_clean", f"uni2_biomedclip_fold{k}", "head.pt")
    if not os.path.exists(f):
        state, best = train_text_head(tr_t, va_t)
        os.makedirs(os.path.dirname(f), exist_ok=True); torch.save(state, f)
        print(f"text head fold {k}: val R@1 {best:.3f} (train pairs {len(tr_t)}, val pairs {len(va_t)})", flush=True)
    text_heads.append(load_mlp(torch.load(f, map_location="cpu")))
    tcga_test.append(one_text_per_image(te_t, np.random.default_rng(k)))
    tcga_notext_test.append(np.flatnonzero((sp.loc[NOTXT.img_idx] == "test").values))  # rows of NOTXT

# ================= (1) HEST: H&E <-> RNA real; virtual RNA (STPath) on the same test spots
M = pd.read_csv(os.path.join(H, "meta.csv"))
hv = pd.read_csv(os.path.join(CACHE, "stpath", "hest_val_meta.csv"))
assert (hv["sample"].values == M["sample"].values).all() and (hv.spot.values == M.spot.values).all()
he_h = np.load(os.path.join(H, "he_uni2.npy"))[:, 0].astype(np.float32)
real_h = torch.from_numpy(np.load(os.path.join(H, "hvg.npy")).astype(np.float32))
virt_h = virtual_rna("hest_val")
with torch.no_grad():
    for k, (trv, te) in enumerate(GroupKFold(K).split(M, groups=M.group)):  # identical folds to 23_bind_gene.py
        grp = np.random.default_rng(k).permutation(sorted(M.group.iloc[trv].unique()))
        va_g = set(grp[: max(1, len(grp) // 10)])
        sp = np.array(["test"] * len(M), dtype=object)
        sp[trv] = np.where(M.group.iloc[trv].isin(va_g), "val", "train")
        s = pd.DataFrame({"spot_idx": M.index, "sample": M["sample"], "group": M.group, "split": sp})
        s.to_csv(os.path.join(SPLITS, "final", f"hest_fold{k}.csv"), index=False)
        log_split("HEST spot-RNA", "group 5-fold", k, s, "group")
        te_saved = torch.load(os.path.join(RUNS, "gene", f"uni2_mlp_fold{k}", "test.pt"), weights_only=False)["idx"]
        assert np.array_equal(np.sort(te_saved), np.sort(te)), "HEST fold mismatch with 23_bind_gene.py"
        zr, zv, zh = nz(rna_heads[k](real_h[te])), nz(rna_heads[k](virt_h[te])), nz(he_h[te])
        add("HEST", "group 5-fold", k, "real", "H&E<->RNA", zh, zr, "trained pair", False)
        for sub, m in [("all test spots", np.ones(len(te), bool)), ("STPath-unseen test spots", hv.hest_bench_heldout.values[te])]:
            if m.sum() >= G100:
                if sub != "all test spots":
                    add(f"HEST ({sub})", "group 5-fold", k, "real", "H&E<->RNA", zh[m], zr[m], "trained pair", False)
                add(f"HEST ({sub})" if sub != "all test spots" else "HEST", "group 5-fold", k, "virtual",
                    "H&E<->RNA(virtual)", zh[m], zv[m], "calibration: virtual replaces real", True)
                add(f"HEST ({sub})" if sub != "all test spots" else "HEST", "group 5-fold", k, "virtual",
                    "RNA<->RNA(virtual)", zr[m], zv[m], "calibration: real vs virtual", False)
print("HEST done", flush=True)

# ================= (2) ORION: H&E <-> mIF real; virtual RNA / synthetic text / virtual mIF on the same test tiles
df = load_all_tiles()
assert (pd.read_csv(os.path.join(CACHE, "stpath", "orion_meta.csv")).key.values == df.image_path.values).all()
he_o = np.load(os.path.join(CACHE, "he_uni2.npy"))[:, 0].astype(np.float32)
rna_o = virtual_rna("orion")
if17 = np.load(os.path.join(CACHE, "if_virtues.npy"))[:, 0].astype(np.float32)
ifs_r = np.load(os.path.join(CACHE, "if_virtues_shared.npy"))[:, 0].astype(np.float32)
ifs_v = np.load(os.path.join(CACHE, "if_virtues_shared_virtual_orion.npy"))[:, 0].astype(np.float32)
gt_o = torch.from_numpy(np.load(os.path.join(CACHE, "orion_gentext_text_biomedclip.npy"))[:, 0].astype(np.float32))


def qm_file(st, k):
    f = os.path.join(CACHE, f"if_virtues_shared_virtual_qm_{st}_fold{k}.npy")
    return np.load(f)[:, 0].astype(np.float32) if os.path.exists(f) else None


qm_o = lambda k, _c={}: _c.setdefault(k, qm_file("orion", k))
# MIPHEI-ViT virtual mIF on the 16-channel ORION panel (30_ --m16). MIPHEI was trained on THIS ORION tile set with the
# official split -> clean only on the official TEST tiles; on patient 5-fold most test tiles were MIPHEI training tiles.
f_m16, f_mo = os.path.join(CACHE, "if_virtues_m16.npy"), os.path.join(CACHE, "if_virtues_m16_virtual_miphei_orion.npy")
m16_r = np.load(f_m16)[:, 0].astype(np.float32) if os.path.exists(f_m16) else None
m16_v = np.load(f_mo)[:, 0].astype(np.float32) if os.path.exists(f_mo) else None
with torch.no_grad():
    for split, folds in [("patient 5-fold", range(K)), ("official", [0])]:
        sdir = "patient_5fold" if split == "patient 5-fold" else "orion_official"
        for k in folds:
            s = pd.read_csv(os.path.join(SPLITS, sdir, f"fold{k}.csv"))
            log_split("ORION H&E-mIF", split, k, s, "patient")
            te = s.tile_id[s.split == "test"].values
            h17, hsh = mif_head("virtues", k, sdir), mif_head("virtues_shared", k, sdir)
            zh, zp, zs = nz(he_o[te]), nz(h17(if17[te])), nz(hsh(ifs_r[te]))
            zr, zt = nz(rna_heads[k](rna_o[te])), nz(text_heads[k](gt_o[te]))  # heads of other datasets: never saw ORION
            A = lambda *a: add("ORION", split, k, *a)
            # mIF = the UNIFIED 10-channel panel shared with GigaTIME everywhere; 17-ch kept only as a reference row
            A("real", "H&E<->mIF (17-ch full ORION panel, reference)", zh, zp, "trained pair", False)
            A("real", "H&E<->mIF", zh, zs, "trained pair", False)
            A("virtual", "mIF<->RNA(virtual)", zs, zr, "emergent: real x virtual", False)
            A("virtual", "mIF<->text(synthetic)", zs, zt, "emergent: real x virtual", False)
            A("virtual", "H&E<->RNA(virtual)", zh, zr, "virtual from this H&E", True)
            A("virtual", "H&E<->text(synthetic)", zh, zt, "virtual from this H&E", True)
            A("virtual", "RNA(virtual)<->text(synthetic)", zr, zt, "emergent: virtual x virtual", True)
            if m16_r is not None and os.path.exists(os.path.join(RUNS, "bench_fm", f"{sdir}_fold{k}_uni2_fm-virtues_m16_fm-geneformer_pathgenclip", "heads.pt")):
                h16 = mif_head("virtues_m16", k, sdir)
                z16 = nz(h16(m16_r[te]))
                A("real", "H&E<->mIF (16-ch MIPHEI panel)", zh, z16, "trained pair", False)
                A("virtual", "mIF(16ch)<->RNA(virtual)", z16, zr, "emergent: real x virtual", False)
                A("virtual", "mIF(16ch)<->text(synthetic)", z16, zt, "emergent: real x virtual", False)
                if m16_v is not None:
                    seen = split != "official"  # 5-fold test tiles were (mostly) MIPHEI training tiles
                    tag = " [MIPHEI-seen tiles]" if seen else ""
                    zm = nz(h16(m16_v[te]))
                    A("virtual", f"H&E<->mIF(MIPHEI){tag}", zh, zm, "calibration: virtual replaces real", True)
                    A("virtual", f"mIF<->mIF(MIPHEI){tag}", z16, zm, "calibration: real vs virtual", False)
                    A("virtual", f"mIF(MIPHEI)<->RNA(virtual){tag}", zm, zr, "emergent: virtual x virtual", True)
                    A("virtual", f"mIF(MIPHEI)<->text(synthetic){tag}", zm, zt, "emergent: virtual x virtual", True)
            for tag, f in [("", ifs_v[te]), (" [qm]", qm_o(k)[te] if split == "patient 5-fold" and qm_o(k) is not None else None)]:
                if f is None:
                    continue  # [qm] = virtual pixels quantile-mapped onto the fold-k TRAIN real distribution (30_ virtual_qm)
                zv = nz(hsh(f))
                A("virtual", f"H&E<->mIF(virtual){tag}", zh, zv, "calibration: virtual replaces real", True)
                A("virtual", f"mIF<->mIF(virtual){tag}", zs, zv, "calibration: real vs virtual", False)
                A("virtual", f"mIF(virtual){tag}<->RNA(virtual)", zv, zr, "emergent: virtual x virtual", True)
                A("virtual", f"mIF(virtual){tag}<->text(synthetic)", zv, zt, "emergent: virtual x virtual", True)
print("ORION done", flush=True)

# ================= (3) TCGA: H&E <-> text real; virtual RNA / virtual mIF on the same test images
rna_t = virtual_rna("tcga_text")
ifv_t = np.load(os.path.join(CACHE, "if_virtues_shared_virtual_tcga_text.npy"))[:, 0].astype(np.float32)
CM = pd.read_csv(os.path.join(CACHE, "tcga_calib_meta.csv"))
syn_t = torch.from_numpy(np.load(os.path.join(CACHE, "tcga_calib_text_biomedclip.npy"))[:, 0].astype(np.float32))
with torch.no_grad():
    for k, (ii, tt) in enumerate(tcga_test):
        zh, zt = nz(IMG[ii, 0]), nz(text_heads[k](TXT[tt]))
        zr = nz(rna_heads[k](rna_t[ii]))
        A = lambda *a: add("TCGA", "patient 5-fold", k, *a)
        A("real", "H&E<->text", zh, zt, "trained pair", False)
        A("virtual", "text<->RNA(virtual)", zt, zr, "emergent: real x virtual", False)
        A("virtual", "H&E<->RNA(virtual)", zh, zr, "virtual from this H&E", True)
        f_mt = os.path.join(CACHE, "if_virtues_m16_virtual_miphei_tcga_text.npy")
        if os.path.exists(f_mt) and os.path.exists(os.path.join(RUNS, "bench_fm", f"patient_5fold_fold{k}_uni2_fm-virtues_m16_fm-geneformer_pathgenclip", "heads.pt")):
            zm = nz(mif_head("virtues_m16", k, "patient_5fold")(np.load(f_mt, mmap_mode="r")[ii, 0].astype(np.float32)))
            A("virtual", "text<->mIF(MIPHEI)", zt, zm, "emergent: real x virtual", False)
            A("virtual", "H&E<->mIF(MIPHEI)", zh, zm, "virtual from this H&E", True)
            A("virtual", "RNA(virtual)<->mIF(MIPHEI)", zr, zm, "emergent: virtual x virtual", True)
        hsh = mif_head("virtues_shared", k, "patient_5fold")  # unified 10-ch panel head, trained on real ORION fold k
        qt = qm_file("tcga_text", k)
        for tag, f in [("", ifv_t), (" [qm]", qt)]:
            if f is None:
                continue
            zv = nz(hsh(f[ii]))
            A("virtual", f"text<->mIF(virtual){tag}", zt, zv, "emergent: real x virtual", False)
            A("virtual", f"H&E<->mIF(virtual){tag}", zh, zv, "virtual from this H&E", True)
            A("virtual", f"RNA(virtual)<->mIF(virtual){tag}", zr, zv, "emergent: virtual x virtual", True)
        # synthetic-caption calibration: test images that have BOTH a real text and a GPT caption
        pos = pd.Series(range(len(ii)), index=ii)
        m = CM.img_idx.isin(ii).values
        if m.sum() >= G100:
            j = pos.loc[CM.img_idx[m]].values
            zs = nz(text_heads[k](syn_t[m]))
            B = lambda *a: add("TCGA (test images with GPT caption)", "patient 5-fold", k, *a)
            B("real", "H&E<->text", zh[j], zt[j], "trained pair", False)
            B("virtual", "H&E<->text(synthetic)", zh[j], zs, "calibration: virtual replaces real", True)
            B("virtual", "text<->text(synthetic)", zt[j], zs, "calibration: real vs virtual", False)
            B("virtual", "text<->RNA(virtual)", zt[j], zr[j], "emergent: real x virtual", False)
            B("virtual", "text(synthetic)<->RNA(virtual)", zs, zr[j], "emergent: virtual x virtual", True)
# images without any real text (after junk removal): GPT few-shot caption, pooled over folds (each image in the test
# split of exactly one fold -> embedded with that fold's heads; all heads map into the same frozen UNI2 space)
f_mt = os.path.join(CACHE, "if_virtues_m16_virtual_miphei_tcga_text.npy")
shot_p = [x["key"][:12] for x in __import__("json").load(open(os.path.join(CACHE, "gen_text", "prompt.json")))["shots"]]
P = {x: [] for x in ["h", "s", "r", "m"]}
with torch.no_grad():
    for k, r in enumerate(tcga_notext_test):
        r = r[~NOTXT.patient.iloc[r].isin(shot_p).values]  # drop few-shot patients
        ii = NOTXT.img_idx.values[r]
        P["h"].append(nz(IMG[ii, 0])); P["s"].append(nz(text_heads[k](TX_N[r]))); P["r"].append(nz(rna_heads[k](rna_t[ii])))
        if os.path.exists(f_mt):
            P["m"].append(nz(mif_head("virtues_m16", k, "patient_5fold")(np.load(f_mt, mmap_mode="r")[ii, 0].astype(np.float32))))
P = {x: torch.cat(v) for x, v in P.items() if v}
B = lambda *a: add("TCGA (test images WITHOUT real text, pooled folds)", "patient 5-fold", "pooled", *a)
B("virtual", "H&E<->text(synthetic)", P["h"], P["s"], "virtual from this H&E", True)
B("virtual", "text(synthetic)<->RNA(virtual)", P["s"], P["r"], "emergent: virtual x virtual", True)
if "m" in P:
    B("virtual", "text(synthetic)<->mIF(MIPHEI)", P["s"], P["m"], "emergent: virtual x virtual", True)
print(f"TCGA no-text images evaluated: {len(P['h'])}")
print("TCGA done", flush=True)

R = pd.DataFrame(rows)
R.to_csv(os.path.join(OUT, "final_long.csv"), index=False)
SP = pd.DataFrame(split_rows)
SP.to_csv(os.path.join(OUT, "final_splits.csv"), index=False)
cols = ["R@1", "R@5", "R@10", "MRR", "pair_acc"]
S = R.groupby(["part", "dataset", "split", "pair", "kind", "circular"], sort=False)[cols + ["n_test"]].agg(["mean", "std"])
S.round(4).to_csv(os.path.join(OUT, "final_summary.csv"))
pd.set_option("display.width", 250); pd.set_option("display.max_colwidth", 50); pd.set_option("display.max_rows", 200)
print(SP.groupby(["dataset", "split", "role"])[["n_items", "n_patient", "n_group"]].mean().round(0).to_string())
for part in ["real", "virtual"]:
    print(f"\n===== {part} (test split, mean over folds) =====")
    print(R[R.part == part].groupby(["dataset", "split", "pair", "kind", "circular"], sort=False)[cols + ["n_test"]].mean().round(3).to_string())