"""Frozen spatial-proteomics foundation-model embeddings of every ORION multiplex-IF tile under the 8 dihedral transforms.
Usage: python 14_extract_if_fm.py virtues [--no_hoechst] [--qc_drop] [--bs 4] [--no_probe]
Output: cache/if_virtues.npy (N=3187, 8, 512) float16 (raw features; row i = tile_id i, axis 1 = common.dihedral d).
        (--no_hoechst / --qc_drop write cache/if_virtues_nohoechst.npy / if_virtues_qcdrop.npy instead.)

VirTues (Wenckstern et al., Nature 2026; github.com/bunnelab/virtues, HF bunnelab/virtues, weights virtues-sp32)
---------------------------------------------------------------------------------------------------------------
Model: MultiplexVirtues with the released base_config (patch 8, dim 512, encoder 'hvhvhvhvhvhvhvhv', RoPE), ESM-2
marker priors. Only the encoder is used (no masking). Input = list of (C,H,W) standardised crops + marker indices.

Tile embedding (pooling): mean over the spatial grid of VirTues *patch summary tokens* (the per-8x8-patch summary
token that attends across all markers; `encoder.forward_list(...).patch_summary_tokens`, shape (H/8, W/8, 512)).
VirTues has no global/CLS token; its own tile/tissue-level tooling (spora-bench SporaVirTuesWrapper.embed_tile /
embed_tissue) returns exactly these patch summary tokens as a bag (then ABMIL), and cell/niche tokens are
area-weighted averages of them (virtues.utils.cell_tokens). A uniform mean over the whole tile = the niche token of
the full 166 um tile. Encoding runs in fp16 autocast, like the spora-bench wrapper.

Marker -> UniProt (ESM-2 esm2_t30_150M_UR50D, mean over residues of layer 30, computed with fair-esm exactly as
virtues/utils/compute_esm_embeddings.py; FASTA from rest.uniprot.org; saved as ext/marker_embeddings/esm2_t30_150M_UR50D/[ID].pt):
  Hoechst P68431 (Histone H3.1)  -- spora flags DNA stains `is_nuclear_marker` with no UniProt ID, so VirTues'
                                    `uniprot_filtered` loading DROPS them (spora-bench VirTues config for lin2023highsubset
                                    = the ORION cohort). spora-io's `replace_nuclear_uniprot_ids=True` instead maps them to
                                    Histone H3 P68431 (used for lin2023high in spora-bench). Default here: keep Hoechst as
                                    P68431 (H3 IMC channels are in VirTues' pretraining vocabulary); --no_hoechst drops it.
  CD31 P16284 (PECAM1) | CD45 P08575 (PTPRC canonical, as in spora examples) | CD68 P34810 | CD4 P01730 | FOXP3 Q9BZS1
  CD8a P01732 | CD45RO P08575-4 (UniProt PTPRC isoform 2 = "CD45R0/CD45RO"; spora uses isoform IDs, e.g. CD45RA P08575-8)
  CD20 P11836 (MS4A1) | PD-L1 Q9NZQ7 (CD274) | CD3e P07766 | CD163 Q86VB7 | E-cadherin P12830 (CDH1) | PD-1 Q15116
  Ki67 P46013 | Pan-CK P02533 (KRT14; the keratin spora/VirTues assign to pan-cytokeratin channels, e.g. cords2024cancer
  "Pan Cytokeratin + Keratin Epithelial") | SMA P62736 (ACTA2)

Preprocessing (mirrors spora-bench configs/datasets/virtues/lin2023highsubset.yaml:
  standardization quantile_clipping_log1p/uq_0.99_image, resolution 1.0 um/px, filter gaussian_blur, use_mean_std,
  disable_quantile_mask=False):
  1. resolution: ORION tiles are 512 px crops of the 0.325 um/px Orion scan resized to 333 px (0.4997 um/px, 166.4 um).
     Resampled to 168x168 px (0.991 um/px; 168 = 21 patches of 8) with antialiased bilinear interpolation.
  2. per-"image" upper quantile: for every slide (spora tissue = one ORION slide) and channel, the 0.99 quantile of all
     resampled pixels of that slide's tiles; x <- log1p(clip(x, 0, uq)).
  3. dataset-level per-channel mean/std of the clipped log1p values over all ORION tiles; x <- (x - mean) / (std + 1e-8).
  4. 3x3 Gaussian blur (sigma 1, reflect padding; torchvision v2.GaussianBlur as in spora-io FilterFactory).
  5. quantile mask: spora drops a (slide, channel) whose upper quantile is 0. On these 8-bit tiles that happens for
     52/697 (slide, channel) pairs (FOXP3 in 25/41 slides, Ki67 6, CD20/CD8a/PD-L1 5, CD31/CD45RO 3) purely because of
     8-bit quantisation of sparse markers, and would give slide-dependent channel sets (a slide/batch confound). Default
     here: uq is floored at 1 grey level and all channels are kept for every tile; --qc_drop reproduces spora's drop.
  6. dihedral d (common.dihedral on the (C,H,W) tensor) applied to the standardised 168 px tile; all 8 encoded.
  The whole 168 px tile is encoded in one pass (21x21 patch grid; RoPE positions so non-128 sizes are supported).

Deviations from the VirTues/spora preprocessing (documented):
  * Input intensities are the 8-bit (0-255) ORION tile exports, not the raw 16-bit Orion counts spora stores; because
    log1p is scale-dependent the clipped-log values are compressed differently (z-scoring per channel partly absorbs it).
  * Quantiles / means / stds come from the 3187 ORION tiles (all pixels, no tissue mask), not whole-slide tissue pixels.
  * VirTues was pretrained on 128x128 crops; tiles here are 168x168 (whole tile, no sliding window).
  * uq floored at 1 grey level instead of spora's per-slide channel drop (step 5).
  * PD-1 in these exports is a near-saturated, almost constant channel (log1p mean 4.65, std 0.16; PD-1_prop ~ 1 for
    most tiles); the per-channel z-scoring therefore amplifies its background texture. Kept as-is (faithful pipeline).
  * Hoechst is kept as Histone H3 (P68431) by default; VirTues' own ORION config drops DNA stains (see --no_hoechst).
  * flash_attn is not installed for this torch build: flash_attn_varlen_qkvpacked_func is replaced by an exact
    PyTorch SDPA implementation over the same variable-length sequences (same math, softmax scale 1/sqrt(d)).
"""
import os
import sys
import time
import types

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")  # less fragmentation on the shared GPU
import numpy as np
import tifffile
import torch
import torch.nn.functional as F

from common import CACHE, ORION, ROOT, dihedral, load_all_tiles

assert sys.argv[1] == "virtues", "usage: python 14_extract_if_fm.py virtues [--no_hoechst] [--qc_drop] [--bs N] [--no_probe]"
NO_HOECHST = "--no_hoechst" in sys.argv
QC_DROP = "--qc_drop" in sys.argv  # spora behaviour: drop (slide, channel) with 0.99-quantile == 0 (default: floor at 1)
BS = int(sys.argv[sys.argv.index("--bs") + 1]) if "--bs" in sys.argv else 4  # tiles per forward (x8 dihedrals = 32 crops, ~4.7 GiB peak)
dev = "cuda"
VIRTUES_REPO = os.path.join(ROOT, "ext", "virtues")
EMB_ROOT = os.path.join(ROOT, "ext", "marker_embeddings")
ESM_NAME = "esm2_t30_150M_UR50D"
EMB_DIR = os.path.join(EMB_ROOT, ESM_NAME)
WEIGHTS = ("virtues-sp32/model.safetensors", "virtues-sp31/model.safetensors")
SIZE, MPP_IN = 168, 0.325 * 512 / 333  # 333 px tile = 512 px @ 0.325 um/px -> 0.4997 um/px; 168 px ~ 0.991 um/px

# channel order of the ORION IF tiles (= the *_prop columns of the dataframes)
MARKERS = [("Hoechst", "P68431"), ("CD31", "P16284"), ("CD45", "P08575"), ("CD68", "P34810"), ("CD4", "P01730"),
           ("FOXP3", "Q9BZS1"), ("CD8a", "P01732"), ("CD45RO", "P08575-4"), ("CD20", "P11836"), ("PD-L1", "Q9NZQ7"),
           ("CD3e", "P07766"), ("CD163", "Q86VB7"), ("E-cadherin", "P12830"), ("PD-1", "Q15116"), ("Ki67", "P46013"),
           ("Pan-CK", "P02533"), ("SMA", "P62736")]


# ----------------------------------------------------------------------------------------------- marker embeddings
def ensure_marker_embeddings(ids):
    """ESM-2 mean-pooled embeddings, identical to virtues/utils/compute_esm_embeddings.py (fair-esm, layer 30)."""
    import requests
    fasta_dir = os.path.join(EMB_ROOT, "fasta")
    os.makedirs(fasta_dir, exist_ok=True)
    os.makedirs(EMB_DIR, exist_ok=True)
    todo = [u for u in ids if not os.path.exists(os.path.join(EMB_DIR, f"{u}.pt"))]
    if not todo:
        return
    import esm
    _load = torch.load  # official fair-esm checkpoints pickle argparse.Namespace -> need weights_only=False
    torch.load = lambda *a, **k: _load(*a, **{**k, "weights_only": False})
    try:
        model, alphabet = esm.pretrained.load_model_and_alphabet(ESM_NAME)
    finally:
        torch.load = _load
    # small one-off job: GPU if the shared card has room, else CPU (the repo recommends CPU for long sequences)
    edev = "cuda" if torch.cuda.is_available() and torch.cuda.mem_get_info()[0] > 4 * 2**30 else "cpu"
    model.eval().to(edev)
    conv = alphabet.get_batch_converter()
    for u in todo:
        fp = os.path.join(fasta_dir, f"{u}.fasta")
        if not os.path.exists(fp):
            r = requests.get(f"https://rest.uniprot.org/uniprotkb/{u}.fasta", timeout=60)
            r.raise_for_status()
            open(fp, "w").write(r.text)
        seq = "".join(l.strip() for l in open(fp) if not l.startswith(">"))
        _, _, tok = conv([(u, seq)])
        n = int((tok != alphabet.padding_idx).sum())
        with torch.no_grad():
            # return_contacts only adds the contact head (repo passes True); representations are identical, and
            # stacking all attention maps for Ki67 (3256 aa) would need 25 GB
            rep = model(tok.to(edev), repr_layers=[30], return_contacts=False)["representations"][30]
        emb = rep[0, 1:n - 1].mean(0).cpu()
        torch.save(emb, os.path.join(EMB_DIR, f"{u}.pt"))
        print(f"  ESM-2 {u}: len {len(seq)} -> {tuple(emb.shape)}", flush=True)
    del model
    torch.cuda.empty_cache()


def check_against_repo_embeddings(ids):
    """The repo ships ESM-2 embeddings for its example dataset; compare on shared UniProt IDs."""
    ref = os.path.join(VIRTUES_REPO, "assets", "example_dataset", "marker_embeddings")
    shared = [u for u in ids if os.path.exists(os.path.join(ref, f"{u}.pt"))]
    if shared:
        cs = [F.cosine_similarity(torch.load(os.path.join(EMB_DIR, f"{u}.pt")).float(),
                                  torch.load(os.path.join(ref, f"{u}.pt")).float(), 0).item() for u in shared]
        mx = max((torch.load(os.path.join(EMB_DIR, f"{u}.pt")) - torch.load(os.path.join(ref, f"{u}.pt"))).abs().max().item()
                 for u in shared)
        print(f"marker embeddings vs repo example set ({len(shared)} shared IDs): min cos {min(cs):.6f}, max|diff| {mx:.2e}")


# ----------------------------------------------------------------------------------------------- flash-attn shim
def install_flash_attn_shim():
    """Exact SDPA replacement for flash_attn_varlen_qkvpacked_func (qkv: (total, 3, H, D), cu_seqlens: (n+1,))."""
    def varlen_qkvpacked(qkv, cu_seqlens, max_seqlen, dropout_p=0.0, softmax_scale=None, causal=False, **_):
        total, _, H, D = qkv.shape
        cu = cu_seqlens.tolist()
        lens = np.diff(cu)
        out = torch.empty(total, H, D, dtype=qkv.dtype, device=qkv.device)
        for L in np.unique(lens):
            if L == 0:
                continue
            starts = np.asarray(cu[:-1])[lens == L]
            if len(starts) == len(lens):  # all sequences same length -> contiguous reshape
                q = qkv.view(len(lens), L, 3, H, D)
                idx = None
            else:
                idx = (torch.as_tensor(starts, device=qkv.device)[:, None] + torch.arange(L, device=qkv.device)).reshape(-1)
                q = qkv[idx].view(len(starts), L, 3, H, D)
            q = q.permute(2, 0, 3, 1, 4)  # 3, n, H, L, D
            o = F.scaled_dot_product_attention(q[0], q[1], q[2], dropout_p=dropout_p, is_causal=causal,
                                               scale=softmax_scale)  # n, H, L, D
            o = o.transpose(1, 2).reshape(-1, H, D)
            if idx is None:
                out = o.contiguous()
            else:
                out[idx] = o
        return out

    pkg = types.ModuleType("flash_attn")
    pkg.__path__ = []
    sub = types.ModuleType("flash_attn.flash_attn_interface")
    sub.flash_attn_varlen_qkvpacked_func = varlen_qkvpacked
    pkg.flash_attn_interface = sub
    sys.modules.setdefault("flash_attn", pkg)
    sys.modules.setdefault("flash_attn.flash_attn_interface", sub)


def load_virtues(marker_ids):
    install_flash_attn_shim()
    if VIRTUES_REPO not in sys.path:
        sys.path.insert(0, VIRTUES_REPO)
    from huggingface_hub import hf_hub_download
    from safetensors.torch import load_file
    from virtues.modules.multiplex_virtues import MultiplexVirtues
    from virtues.utils.utils import load_marker_embedding_dict, load_marker_embeddings

    emb = load_marker_embeddings(EMB_DIR)
    emb_dict = load_marker_embedding_dict(EMB_DIR)
    model = MultiplexVirtues(prior_bias_embeddings=emb)  # defaults == configs/base_config.yaml (released model)
    for w in WEIGHTS:
        try:
            sd = load_file(hf_hub_download("bunnelab/virtues", w))
            model.load_state_dict(sd, strict=True)
            print("loaded", w)
            break
        except Exception as e:  # fall back to sp31
            print("could not load", w, repr(e)[:200])
    else:
        raise RuntimeError("no VirTues weights could be loaded")
    model.eval().to(dev)
    midx = torch.tensor([emb_dict[u] for u in marker_ids], dtype=torch.long)
    return model, midx


# ----------------------------------------------------------------------------------------------- data
def load_resampled(df):
    """Raw uint8 tiles -> (N,17,SIZE,SIZE) float16 at ~1 um/px, plus raw per-channel mean intensity (baseline)."""
    N = len(df)
    X = np.empty((N, 17, SIZE, SIZE), np.float16)
    raw_mean = np.empty((N, 17), np.float32)
    paths = [os.path.join(ORION, p) for p in df.target_path]
    chunk = 64
    for s in range(0, N, chunk):
        t = np.stack([tifffile.imread(p) for p in paths[s:s + chunk]])  # n,333,333,17 uint8
        x = torch.from_numpy(t).to(dev).permute(0, 3, 1, 2).float()
        raw_mean[s:s + chunk] = x.mean((2, 3)).cpu().numpy()
        X[s:s + chunk] = F.interpolate(x, size=(SIZE, SIZE), mode="bilinear", antialias=True,
                                       align_corners=False).clamp(min=0).half().cpu().numpy()
        if s % 640 == 0:
            print(f"  read {s}/{N}", flush=True)
    return X, raw_mean


def standardize(X, slide):
    """quantile_clipping_log1p / uq_0.99_image + global mean/std + gaussian blur (in place, float16)."""
    from torchvision.transforms.v2.functional import gaussian_blur
    N, C = X.shape[:2]
    slides = np.unique(slide)
    uq = np.zeros((len(slides), C), np.float32)
    n_zero = np.zeros((len(slides), C), bool)
    s1, s2, cnt = np.zeros(C), np.zeros(C), np.zeros(C)
    for k, sl in enumerate(slides):
        idx = np.where(slide == sl)[0]
        x = torch.from_numpy(X[idx]).to(dev).float().transpose(0, 1).reshape(C, -1)  # C, n*H*W
        # torch.quantile is limited to 2**24 elements: very large slides use a fixed random pixel subsample
        qs = x if x.shape[1] <= 2 ** 24 else x[:, torch.randperm(x.shape[1], generator=torch.Generator().manual_seed(0))[: 2 ** 24].to(x.device)]
        uq[k] = np.array([torch.quantile(qs[c], 0.99).item() for c in range(C)])
        n_zero[k] = uq[k] <= 0
        if not QC_DROP:  # 8-bit tiles: sparse markers (FOXP3, ...) often have uq == 0 -> floor at one grey level
            uq[k] = np.maximum(uq[k], 1.0)
        u = torch.from_numpy(uq[k]).to(dev)[:, None]
        v = torch.log1p(torch.minimum(x, u))
        ok = (uq[k] > 0)
        s1[ok] += v.sum(1).double().cpu().numpy()[ok]
        s2[ok] += (v.double() ** 2).sum(1).cpu().numpy()[ok]
        cnt[ok] += v.shape[1]
    mean = s1 / cnt
    std = np.sqrt(s2 / cnt - mean ** 2)
    valid = uq > 0  # spora quantile mask: drop (slide, channel) with zero upper quantile
    slide_k = {sl: k for k, sl in enumerate(slides)}
    m_t = torch.tensor(mean, dtype=torch.float32, device=dev).view(1, C, 1, 1)
    s_t = torch.tensor(std, dtype=torch.float32, device=dev).view(1, C, 1, 1)
    for s in range(0, N, 128):
        ks = [slide_k[sl] for sl in slide[s:s + 128]]
        u = torch.from_numpy(uq[ks]).to(dev)[:, :, None, None]
        x = torch.from_numpy(X[s:s + 128]).to(dev).float()
        x = (torch.log1p(torch.minimum(x, u)) - m_t) / (s_t + 1e-8)
        x = gaussian_blur(x, kernel_size=[3, 3], sigma=[1.0, 1.0])
        X[s:s + 128] = x.half().cpu().numpy()
    return X, uq, mean, std, valid, slide_k, n_zero


# ----------------------------------------------------------------------------------------------- sanity checks
def sanity(out, raw_mean, markers, groups):
    from scipy.stats import spearmanr
    from sklearn.linear_model import RidgeCV
    from sklearn.model_selection import GroupKFold, KFold, cross_val_predict
    from sklearn.pipeline import make_pipeline
    from sklearn.preprocessing import StandardScaler
    o = out.astype(np.float32)
    print("\n=== sanity checks ===")
    print("finite:", bool(np.isfinite(o).all()), "| shape", out.shape, out.dtype)
    z = o[:, 0]
    zn = z / np.linalg.norm(z, axis=1, keepdims=True)
    rng = np.random.default_rng(0)
    s = rng.choice(len(z), min(1000, len(z)), replace=False)
    cos = zn[s] @ zn[s].T
    off = cos[~np.eye(len(s), dtype=bool)]
    zc = z - z.mean(0)
    zcn = zc / np.linalg.norm(zc, axis=1, keepdims=True)
    offc = (zcn[s] @ zcn[s].T)[~np.eye(len(s), dtype=bool)]
    sd = z.std(0)
    ev = np.linalg.eigvalsh(np.cov(z.T))[::-1]
    print(f"mean pairwise cosine (raw) {off.mean():.4f} (min {off.min():.4f}) | after centering {offc.mean():.4f}")
    print(f"per-dim std: min {sd.min():.4f} median {np.median(sd):.4f} max {sd.max():.4f}; "
          f"dims with std<1e-3: {(sd < 1e-3).sum()}; PCA participation ratio {ev.sum() ** 2 / (ev ** 2).sum():.1f}; "
          f"top-10 PCs explain {ev[:10].sum() / ev.sum():.3f}")
    on = o / np.linalg.norm(o, axis=2, keepdims=True)
    same = np.mean([(on[:, 0] * on[:, d]).sum(1).mean() for d in range(1, 8)])
    samec = np.mean([(zcn * ((o[:, d] - z.mean(0)) / np.linalg.norm(o[:, d] - z.mean(0), axis=1, keepdims=True))).sum(1).mean()
                     for d in range(1, 8)])
    print(f"dihedral consistency: mean cos(d0, d1..7) same tile {same:.4f} (centered {samec:.4f}) "
          f"vs random tile pairs {off.mean():.4f} (centered {offc.mean():.4f})")

    Y = markers[:, :17]
    names = [m for m, _ in MARKERS]

    def probe(Xf, cv, grp=None):
        pipe = make_pipeline(StandardScaler(), RidgeCV(alphas=np.logspace(-2, 5, 15)))
        P = cross_val_predict(pipe, Xf, Y, cv=cv, groups=grp)
        r = np.array([spearmanr(P[:, j], Y[:, j])[0] for j in range(Y.shape[1])])
        return np.nan_to_num(r)

    feats = {"VirTues (d0)": z, "VirTues (mean 8 dihedral)": o.mean(1), "baseline: IF mean intensity": raw_mean,
             "baseline: log1p IF mean intensity": np.log1p(raw_mean)}
    res = {}
    for cvname, cv, grp in [("5-fold tiles", KFold(5, shuffle=True, random_state=0), None),
                            ("5-fold patient-grouped", GroupKFold(5), groups)]:
        for fn, Xf in feats.items():
            r = probe(Xf, cv, grp)
            res[(cvname, fn)] = r
            print(f"ridge probe [{cvname:22s}] {fn:34s} mean Spearman over 17 *_prop = {r.mean():.4f}")
    print("per-marker Spearman (5-fold tiles): VirTues d0 | IF mean intensity")
    a, b = res[("5-fold tiles", "VirTues (d0)")], res[("5-fold tiles", "baseline: IF mean intensity")]
    print("  " + "  ".join(f"{n}:{x:.2f}|{y:.2f}" for n, x, y in zip(names, a, b)))
    return res


# ----------------------------------------------------------------------------------------------- main
if __name__ == "__main__":
    t0 = time.time()
    df = load_all_tiles()
    N = len(df)
    keep_ch = [c for c in range(17) if not (NO_HOECHST and MARKERS[c][0] == "Hoechst")]
    ids = [MARKERS[c][1] for c in range(17)]
    print("marker -> UniProt:", ", ".join(f"{m}={u}" for m, u in MARKERS), "| using channels:", len(keep_ch))
    ensure_marker_embeddings(sorted(set(ids)))
    check_against_repo_embeddings(sorted(set(ids)))
    model, midx_all = load_virtues(ids)
    print(f"[{time.time() - t0:.0f}s] model ready; reading tiles (input {MPP_IN:.4f} um/px -> {SIZE}px = "
          f"{MPP_IN * 333 / SIZE:.3f} um/px)", flush=True)

    X, raw_mean = load_resampled(df)
    slide = df.in_slide_name.values
    X, uq, mean, std, valid, slide_k, n_zero = standardize(X, slide)
    print(f"[{time.time() - t0:.0f}s] standardized. global log1p mean {np.round(mean, 2).tolist()}\n  std {np.round(std, 2).tolist()}")
    print(f"  (slide, channel) pairs with 0.99-quantile == 0: {n_zero.sum()} of {n_zero.size} ("
          + ", ".join(f"{MARKERS[c][0]}x{n_zero[:, c].sum()}" for c in range(17) if n_zero[:, c].any())
          + (") -> dropped (spora quantile mask, --qc_drop)" if QC_DROP else ") -> uq floored at 1 grey level, channel kept"))

    # per-tile channel set (quantile mask x optional Hoechst removal) -> batch tiles with identical channel sets
    tile_valid = valid[[slide_k[s] for s in slide]]
    tile_valid[:, [c for c in range(17) if c not in keep_ch]] = False
    sig = [tuple(np.where(v)[0]) for v in tile_valid]
    groups = {}
    for i, g in enumerate(sig):
        groups.setdefault(g, []).append(i)

    out = np.zeros((N, 8, 512), np.float16)
    torch.cuda.reset_peak_memory_stats()
    done = 0
    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16):
        for g, idxs in groups.items():
            ch = torch.tensor(g, dtype=torch.long)
            midx = midx_all[ch].to(dev)
            for s in range(0, len(idxs), BS):
                b = idxs[s:s + BS]
                x = torch.from_numpy(X[b]).to(dev).float()[:, ch]  # B,C,H,W
                xd = [dihedral(x, d) for d in range(8)]
                crops = [xd[d][i] for d in range(8) for i in range(len(b))]  # order: d-major
                ps = model.encoder.forward_list(crops, [midx] * len(crops)).patch_summary_tokens  # list of (21,21,512)
                e = torch.stack([p.float().mean((0, 1)) for p in ps]).view(8, len(b), -1).transpose(0, 1)
                out[b] = e.cpu().numpy().astype(np.float16)
                done += len(b)
                if (done // BS) % 50 == 0:
                    print(f"  encoded {done}/{N} [{time.time() - t0:.0f}s] peak GPU "
                          f"{torch.cuda.max_memory_allocated() / 2**30:.2f} GiB", flush=True)
    fn = os.path.join(CACHE, "if_virtues" + ("_nohoechst" if NO_HOECHST else "") + ("_qcdrop" if QC_DROP else "") + ".npy")
    np.save(fn, out)
    print(f"saved {fn} {out.shape} {out.dtype} finite: {np.isfinite(out).all()} | channel-set groups: {len(groups)} | "
          f"peak GPU alloc {torch.cuda.max_memory_allocated() / 2**30:.2f} GiB | encode+prep time {time.time() - t0:.0f}s")

    if "--no_probe" not in sys.argv:
        sanity(out, raw_mean, np.load(os.path.join(CACHE, "markers.npy")), df.patient.values)
    print(f"total {time.time() - t0:.0f}s")