"""STPath (Huang et al. 2025) zero-shot spatial gene-expression prediction from H&E GigaPath features.

Source   code    https://github.com/Graph-and-Geometric-Learning/STPath  (commit 3b10e6d, no LICENSE file)
         weights https://huggingface.co/tlhuang/STPath  stfm.pth  (no licence declared; not gated)
         local copy: ext/stpath/{repo,weights}   -- treated as untrusted: state dict loaded with weights_only=True,
         and only the pure-torch model modules (model.py, spatial_transformer.py, fa.py, nn_utils) are imported.
Model    STFM spatial transformer: per-spot input = Linear(GigaPath tile feature 1536) + masked gene token
         + organ token + tech token; attention over all spots of a slide with a 2D frame-averaged distance bias.
         Coordinates are shifted to 0 and rescaled to [0,100] per slide (official STPathInference._normalize_coords).
         Output: 38,984 Ensembl-gene vector per spot on the training scale log1p(raw UMI count) (NOT library-normalised;
         training used normalize_method='log1pv2' = sc.pp.log1p on raw counts).
Encoder  The released checkpoint was trained on prov-gigapath tile features ONLY (official pipeline: 224-px HEST patches
         @ 0.5 um/px, CenterCrop 224, ImageNet norm).  prov-gigapath is gated: --extract needs HF_TOKEN in the env
         (pass it as an env var only; this script never reads, prints or stores it).
         Our features = mean of GigaPath(fp16) over the centre + 4 corner 224-px crops at ~0.5 um/px
         (centre-only saved too as <set>_gigapath_center.npy):
           ORION  333-px tiles (512-px @ 0.325 um/px resized -> ~0.49 um/px), cropped at native size.
           TCGA   672-px PathGen patches = read_region level 0; visually 40x (RBC ~28-30 px -> ~0.25 um/px),
                  so resized 672 -> 336 px (~0.5 um/px) before cropping.  (20x-scanned slides, if any, would be off 2x.)

Sets
  int2_check  repo example (HEST INT2, CCRCC, full slide 2541 spots) -> reproduces README mean Pearson 0.1562;
              also INT2 in random 26-spot groups (sparse context, like our 1% HEST subsample) and single spots.
  hest_val    cache/hest (9484 Visium spots, 1% subsample of 421 HEST samples, precomputed GigaPath features),
              one STPath pass per sample (all its spots as context), organ = HEST organ, tech = Visium.
              Samples in HEST-bench were excluded from STPath pretraining -> reported separately as held-out.
  orion       ORION CRC tiles (common.load_all_tiles, 3187 tiles / 41 slides), one pass per slide with coords from
              the filename x_y (level-0 px), organ = Bowel, tech = Visium.  + protein sanity check
              (orion_protein_check.csv: predicted RNA vs <marker>_count/n_nuclei, tiles with >= 20 nuclei).
  tcga_text   cache/allimg_meta.csv (6144 patches, 4088 WSIs -> mostly 1 patch = no spatial context), one pass per
              WSI, coords from key x_y; organ from the TCGA project (public GDC files API, cached in
              tcga_text_projects.csv) via PROJECT2ORGAN, unmapped -> 'Others'; tech = Visium.

Outputs cache/stpath/<set>_pred.npy (N,G) float16, native STPath scale log1p(raw count); <set>_genes.csv; <set>_meta.csv
        hest_val additionally: hest_val_metrics.csv (per-gene Pearson/Spearman), hest_val_within_sample.csv.
Run     .venv/bin/python 26_stpath_infer.py --set int2_check|hest_val
        HF_TOKEN=... .venv/bin/python 26_stpath_infer.py --set orion|tcga_text --extract   # GigaPath features
        .venv/bin/python 26_stpath_infer.py --set orion|tcga_text [--feats F.npy]          # STPath prediction
"""
import argparse
import json
import os
import pickle
import re
import sys
import weakref

import numpy as np
import pandas as pd
import torch

ROOT = os.path.dirname(os.path.abspath(__file__))
EXT = os.path.join(ROOT, "ext", "stpath")
REPO = os.path.join(EXT, "repo")
WEIGHTS = os.path.join(EXT, "weights", "stfm.pth")
VOC = os.path.join(REPO, "utils_data", "symbol2ensembl.json")
CACHE = os.path.join(ROOT, "cache")
OUT = os.path.join(CACHE, "stpath")
HEST = os.path.join(ROOT, "..", "HEST")
MARKERS = ["PTPRC", "EPCAM", "KRT8", "ACTA2", "MKI67", "PECAM1", "CD68", "CD3E", "CDH1", "COL1A1"]
HEST_BENCH = {  # MahmoodLab/hest-bench adata ids; STPath excluded these from pretraining (train.py)
    "INT%d" % i for i in range(1, 25)} | {"TENX111", "TENX147", "TENX148", "TENX149", "NCBI642", "NCBI643",
    "NCBI783", "NCBI785", "TENX95", "TENX99", "TENX118", "TENX141", "NCBI681", "NCBI682", "NCBI683", "NCBI684",
    "TENX116", "TENX126", "TENX140", "ZEN36", "ZEN40", "ZEN48", "ZEN49", "TENX115", "TENX117"} | {
    "MEND%d" % i for i in list(range(139, 155)) + list(range(156, 163))}


# ----------------------------------------------------------------------------------------------- model
class STPath:
    def __init__(self, device):
        sys.path.insert(0, REPO)
        from stpath.model.model import STFM
        from stpath.model.nn_utils.config import ModelConfig
        import stpath.utils.constants as C

        sym2ens = json.load(open(VOC))
        ens = sorted(set(sym2ens.values()))
        self.gene2id = {g: i + 2 for i, g in enumerate(ens)}  # 0 pad, 1 mask (GeneExpTokenizer)
        self.sym2ens = sym2ens
        self.n_genes = max(self.gene2id.values()) + 1
        self.C = C
        cfg = ModelConfig.get_default_config()
        cfg.feature_dim, cfg.activation, cfg.backbone = 1536, "gelu", "spatial_transformer"
        cfg.n_genes, cfg.n_tech, cfg.n_species, cfg.n_organs = self.n_genes, len(C.tech_voc), len(C.specie_voc), len(C.organ_voc)
        self.model = STFM(cfg)
        self.model.load_state_dict(torch.load(WEIGHTS, map_location="cpu", weights_only=True))
        self.model.eval().to(device)
        self.device = device

    def organ_id(self, organ):  # IDTokenizer.encode(align_first=True); None -> "Others" as in STPathInference
        organ = "Others" if organ is None or (isinstance(organ, float) and np.isnan(organ)) else organ
        return self.C.organ_voc.index(self.C.organ_align_mapping.get(organ, "<pad>"))

    def tech_id(self, tech):
        return 0 if tech is None else self.C.tech_voc.index(self.C.tech_align_mapping.get(tech, "<pad>"))

    def gene_cols(self, symbols):
        """token column for each symbol (-1 if not in the STPath vocabulary)"""
        return np.array([self.gene2id.get(self.sym2ens.get(s), -1) for s in symbols])

    @torch.no_grad()
    def predict(self, feats, coords, organ, tech, cols):
        """feats (n,1536), coords (n,2) one slide -> (n, len(cols)) log1p(count)-scale, NaN for cols==-1"""
        n = len(feats)
        c = torch.as_tensor(np.asarray(coords, np.float64), dtype=torch.float32)
        c = c - c.min(0).values
        rng = c.max(0).values
        c = torch.where(rng > 0, c / rng.clamp(min=1e-8) * 100, torch.zeros_like(c))  # rescale_coords; guard 1-spot/collinear
        x = torch.as_tensor(np.asarray(feats, np.float32), device=self.device)
        ge = torch.zeros(n, self.n_genes, device=self.device)
        ge[:, 1] = 1  # all spots masked (pure image -> expression)
        p = self.model.prediction_head(
            img_tokens=x, coords=c.to(self.device), ge_tokens=ge,
            batch_idx=torch.zeros(n, dtype=torch.long, device=self.device),
            tech_tokens=torch.full((n,), self.tech_id(tech), dtype=torch.long, device=self.device),
            organ_tokens=torch.full((n,), self.organ_id(organ), dtype=torch.long, device=self.device))
        out = p[:, torch.as_tensor(np.clip(cols, 0, None), device=self.device)].float().cpu().numpy()
        out[:, cols < 0] = np.nan
        # library size of the predicted profile (all vocab genes) for optional CP10k re-normalisation
        lib = torch.expm1(p[:, 2:].clamp(min=0)).sum(1).cpu().numpy()
        return out, lib


# ----------------------------------------------------------------------------------------------- metrics
def corr_cols(a, b):
    a = a - a.mean(0); b = b - b.mean(0)
    return (a * b).sum(0) / np.sqrt((a * a).sum(0) * (b * b).sum(0) + 1e-20)


def rank_cols(a):
    from scipy.stats import rankdata
    return np.apply_along_axis(rankdata, 0, a)


def gene_metrics(pred, truth):
    return corr_cols(pred, truth), corr_cols(rank_cols(pred), rank_cols(truth))


# ----------------------------------------------------------------------------------------------- sets
def run_int2(m):
    import anndata as ad
    import h5py
    ex = os.path.join(REPO, "example_data")
    hvg = json.load(open(os.path.join(ex, "var_50genes.json")))["genes"]
    with h5py.File(os.path.join(ex, "INT2.h5")) as f:
        E, Cc, bc = f["embeddings"][:], f["coords"][:], f["barcodes"][:].flatten().astype(str)
    a = ad.read_h5ad(os.path.join(ex, "INT2.h5ad"))[bc, :]
    gt = np.log1p(a[:, hvg].X.toarray())
    cols = m.gene_cols(hvg)
    p, _ = m.predict(E, Cc, "Kidney", "Visium", cols)
    r = corr_cols(p, gt)
    print(f"INT2 full slide ({len(E)} spots): mean Pearson 50 HVG = {np.nanmean(r):.4f} (README: 0.1562)")
    rs = []
    for seed in range(5):  # sparse context like our 1% HEST subsample: 26 random spots per pass
        idx = np.random.default_rng(seed).permutation(len(E))
        pp = np.zeros_like(p)
        for ch in np.array_split(idx, len(E) // 26):
            pp[ch] = m.predict(E[ch], Cc[ch], "Kidney", "Visium", cols)[0]
        rs.append(np.nanmean(corr_cols(pp, gt)))
    print(f"INT2 sparse context (random 26-spot groups, 5 seeds): mean Pearson = {np.mean(rs):.4f} +- {np.std(rs):.4f}")
    p1 = np.concatenate([m.predict(E[i:i + 1], Cc[i:i + 1], "Kidney", "Visium", cols)[0] for i in range(len(E))])
    print(f"INT2 no context (each spot alone): mean Pearson = {np.nanmean(corr_cols(p1, gt)):.4f}")


def hest_coords():
    from anndata._core import file_backing

    def _ss(self, state):  # same unpickle patch as 21_hest_prep.py
        self.__dict__ = state.copy()
        ref = state.get("_adata_ref", state.get("_adata"))
        self.__dict__["_adata_ref"] = weakref.ref(ref) if ref is not None else (lambda: None)

    file_backing.AnnDataFileManager.__setstate__ = _ss
    A = pickle.load(open(os.path.join(HEST, "visium_humandisease_adata.pickle"), "rb"))
    H = pd.read_csv(os.path.join(HEST, "HEST_v1_1_0.csv"))
    H = H[(H.st_technology == "Visium") & (H.species == "Homo sapiens")].reset_index(drop=True)
    assert len(H) == len(A)
    xy = {}
    for a, sid in zip(A, H.id):
        for b, c in zip(a.obs_names, np.asarray(a.obsm["spatial"], np.float64)):
            xy[(sid, b)] = c
    return xy


def run_hest(m):
    import scipy.sparse as sp
    M = pd.read_csv(os.path.join(CACHE, "hest", "meta.csv"))
    genes = pd.read_csv(os.path.join(CACHE, "hest", "genes.csv")).gene.tolist()
    X = sp.load_npz(os.path.join(CACHE, "hest", "counts.npz")).tocsr()
    F = np.load(os.path.join(CACHE, "hest", "he_gigapath.npy"))[:, 0].astype(np.float32)
    xy = hest_coords()
    coords = np.stack([xy[(s, b)] for s, b in zip(M["sample"], M.spot)])
    cols = m.gene_cols(genes)
    keep = cols >= 0
    print(f"HEST genes {len(genes)}, in STPath vocab {keep.sum()}")
    genes_k, cols_k = [g for g, k in zip(genes, keep) if k], cols[keep]
    pred = np.zeros((len(M), len(cols_k)), np.float32)
    lib = np.zeros(len(M), np.float32)
    for sid, idx in M.groupby("sample").indices.items():
        pred[idx], lib[idx] = m.predict(F[idx], coords[idx], M.organ.iloc[idx[0]], "Visium", cols_k)
    os.makedirs(OUT, exist_ok=True)
    np.save(os.path.join(OUT, "hest_val_pred.npy"), pred.astype(np.float16))
    pd.Series(genes_k).to_csv(os.path.join(OUT, "hest_val_genes.csv"), index=False, header=["gene"])
    mo = M[["idx", "sample", "spot", "organ", "oncotree"]].copy()
    mo["x"], mo["y"] = coords[:, 0], coords[:, 1]
    mo["hest_bench_heldout"] = mo["sample"].isin(HEST_BENCH)
    mo["pred_lib_expm1"] = lib
    mo["status"] = "ok"
    mo.to_csv(os.path.join(OUT, "hest_val_meta.csv"), index=False)

    # ---- evaluation: truth = log1p(CP10k) of measured counts (and log1p raw = STPath's training target)
    Xk = X[:, np.where(keep)[0]]
    tot = np.asarray(X.sum(1)).ravel()
    L = sp.diags(1e4 / tot) @ Xk
    L.data = np.log1p(L.data)
    frac = np.asarray((Xk > 0).mean(0)).ravel()
    mu = np.asarray(L.mean(0)).ravel()
    var = np.asarray(L.multiply(L).mean(0)).ravel() - mu ** 2
    var[frac < 0.01] = -1
    top50 = list(np.argsort(-var)[:50])
    gi = {g: i for i, g in enumerate(genes_k)}
    mk = [gi[g] for g in MARKERS if g in gi]
    sel = top50 + [i for i in mk if i not in top50]
    T_cp = L[:, sel].toarray()
    T_raw = np.log1p(Xk[:, sel].toarray())
    P = pred[:, sel]
    P_cp = np.log1p(np.expm1(np.clip(P, 0, None)) / lib[:, None] * 1e4)  # predicted profile re-normalised to CP10k
    held = mo.hest_bench_heldout.values
    rows = []
    for subset, msk in [("all", np.ones(len(M), bool)), ("heldout_hestbench", held), ("seen_pretrain", ~held)]:
        for tname, T, PP in [("log1p_cp10k", T_cp, P), ("log1p_raw", T_raw, P), ("log1p_cp10k_vs_predcp10k", T_cp, P_cp)]:
            pr, sr = gene_metrics(PP[msk], T[msk])
            for j, g in enumerate(sel):
                rows.append({"subset": subset, "truth": tname, "gene": genes_k[g], "is_top50_hvg": g in top50,
                             "is_marker": genes_k[g] in MARKERS, "pearson": pr[j], "spearman": sr[j],
                             "n_spots": int(msk.sum())})
    R = pd.DataFrame(rows)
    R.to_csv(os.path.join(OUT, "hest_val_metrics.csv"), index=False)
    # within-sample correlations (spatial-variation within a slide, the HEST-bench style), samples with >=30 spots
    ws = []
    for sid, idx in M.groupby("sample").indices.items():
        if len(idx) >= 30:
            pr, _ = gene_metrics(P[idx][:, :50], T_cp[idx][:, :50])
            ws.append({"sample": sid, "heldout": sid in HEST_BENCH, "n": len(idx), "mean_pearson_top50": np.nanmean(pr)})
    W = pd.DataFrame(ws)
    W.to_csv(os.path.join(OUT, "hest_val_within_sample.csv"), index=False)
    print(f"\nspots {len(M)}, samples {M['sample'].nunique()}; held-out (HEST-bench) spots {held.sum()} "
          f"in {M[held]['sample'].nunique()} samples")
    s = R.groupby(["subset", "truth", "is_top50_hvg"])[["pearson", "spearman"]].agg(["mean", "median"])
    print(s.round(3).to_string())
    print("\nmarkers, truth log1p CP10k (pooled across spots):")
    print(R[R.is_marker & (R.truth == "log1p_cp10k")].pivot(index="gene", columns="subset", values=["pearson", "spearman"]).round(3).to_string())
    print("\nwithin-sample mean Pearson top-50 HVG (samples >=30 spots):")
    print(W.groupby("heldout").mean_pearson_top50.describe().round(3).to_string())


def parse_xy(s):
    mm = re.search(r"_(\d+)_(\d+)(?:_0_512_512)?(?:\.\w+)?$", s)
    return (float(mm.group(1)), float(mm.group(2))) if mm else (np.nan, np.nan)


def target_paths(name):
    if name == "orion":
        sys.path.insert(0, ROOT)
        from common import load_all_tiles, ORION
        return [os.path.join(ORION, p) for p in load_all_tiles().image_path], 333   # ~0.49 um/px already
    d = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv"))
    return d.image_path.tolist(), 336  # PathGen read_region level 0 = 40x (~0.25 um/px): 672 -> 336 = ~0.5 um/px


def five_crops(img, size=224):
    W, H = img.size
    offs = [((W - size) // 2, (H - size) // 2), (0, 0), (W - size, 0), (0, H - size), (W - size, H - size)]
    return [img.crop((x, y, x + size, y + size)) for x, y in offs]  # centre first


def extract_gigapath(name, device, bs=32):
    """GigaPath tile features (fp16) = mean over centre + 4 corner 224-px crops at ~0.5 um/px; centre-only also saved.
    Needs HF_TOKEN in the environment (prov-gigapath is gated); the token is never read/stored by this script."""
    import timm
    from PIL import Image, PngImagePlugin
    from torchvision import transforms
    PngImagePlugin.MAX_TEXT_CHUNK = 64 * 1024 * 1024
    enc = timm.create_model("hf_hub:prov-gigapath/prov-gigapath", pretrained=True).half().eval().to(device)
    tf = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])
    paths, side = target_paths(name)
    feats = np.full((len(paths), 5, 1536), np.nan, np.float32)
    status = ["ok"] * len(paths)
    buf, rows = [], []

    def flush():
        if not buf:
            return
        with torch.inference_mode():
            f = enc(torch.stack(buf).half().to(device)).float().cpu().numpy()
        feats[np.array(rows)] = f.reshape(-1, 5, 1536)
        buf.clear(); rows.clear()

    for i, p in enumerate(paths):
        try:
            im = Image.open(p).convert("RGB")
            if im.size != (side, side):
                im = im.resize((side, side), Image.BICUBIC)
            buf.extend(tf(c) for c in five_crops(im)); rows.append(i)
        except Exception as e:
            status[i] = f"unreadable:{type(e).__name__}"
            print(f"[{name}] skip {p}: {e}", flush=True)
        if len(rows) * 5 >= bs:
            flush()
        if i % 500 == 0:
            print(f"[{name}] {i}/{len(paths)}", flush=True)
    flush()
    os.makedirs(OUT, exist_ok=True)
    np.save(os.path.join(OUT, f"{name}_gigapath.npy"), feats.mean(1))
    np.save(os.path.join(OUT, f"{name}_gigapath_center.npy"), feats[:, 0])
    pd.DataFrame({"path": paths, "feat_status": status}).to_csv(os.path.join(OUT, f"{name}_gigapath_status.csv"), index_label="row")
    print(f"[{name}] features {feats.shape[0]} rows, failed {sum(s != 'ok' for s in status)}; "
          f"non-finite rows {int((~np.isfinite(feats.mean(1)).all(1)).sum())}")


def tcga_projects(meta):
    """TCGA project per WSI from the public GDC files API (file_name = <wsi>.svs); cached in cache/stpath."""
    fp = os.path.join(OUT, "tcga_text_projects.csv")
    if not os.path.exists(fp):
        import requests
        names = sorted(set(meta.slide + ".svs"))
        rows = []
        for i in range(0, len(names), 500):
            q = {"filters": {"op": "in", "content": {"field": "file_name", "value": names[i:i + 500]}},
                 "fields": "file_name,cases.project.project_id", "format": "JSON", "size": 1000}
            r = requests.post("https://api.gdc.cancer.gov/files", json=q, timeout=120).json()
            for h in r["data"]["hits"]:
                rows.append({"slide": h["file_name"][:-4], "project": h["cases"][0]["project"]["project_id"]})
        pd.DataFrame(rows).drop_duplicates("slide").to_csv(fp, index=False)
    return dict(pd.read_csv(fp).values)


# TCGA project -> STPath organ name (keys of constants.organ_align_mapping; lower-case ones are the STimage-1K4M
# spellings that align to their own tokens, e.g. 'ovary'->'Ovary'); unlisted projects -> 'Others'
PROJECT2ORGAN = {"BRCA": "Breast", "COAD": "Bowel", "READ": "Bowel", "LUAD": "Lung", "LUSC": "Lung",
                 "KIRC": "Kidney", "KIRP": "Kidney", "KICH": "Kidney", "LIHC": "Liver", "PAAD": "Pancreas",
                 "PRAD": "Prostate", "SKCM": "Skin", "UCEC": "Uterus", "UCS": "Uterus", "OV": "ovary",
                 "GBM": "glioblastoma", "LGG": "glioma", "STAD": "stomach", "HNSC": "mouth", "UVM": "Eye"}


def run_target(m, name, feats_path):
    if name == "orion":
        sys.path.insert(0, ROOT)
        from common import load_all_tiles
        d = load_all_tiles()
        meta = pd.DataFrame({"key": d.image_path, "path": d.image_path, "slide": d.in_slide_name})
        meta["x"], meta["y"] = zip(*[parse_xy(os.path.basename(p)) for p in d.image_path])
        meta["organ"] = "Bowel"
    else:
        d = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv"))
        meta = pd.DataFrame({"key": d.key, "path": d.image_path, "slide": d.wsi_id})
        meta["x"], meta["y"] = zip(*[parse_xy(k) for k in d.key])
        proj = tcga_projects(meta)  # TCGA project from GDC; organ token mapped, unknown -> 'Others'
        meta["project"] = meta.slide.map(proj)
        meta["organ"] = [PROJECT2ORGAN.get(str(p).replace("TCGA-", ""), "Others") for p in meta.project]
    feats_path = feats_path or os.path.join(OUT, f"{name}_gigapath.npy")
    if not os.path.exists(feats_path):
        sys.exit(f"[{name}] BLOCKED: STPath's released checkpoint needs prov-gigapath tile features (1536-d) and "
                 f"prov-gigapath is gated for this HF account. Provide --feats <GigaPath features (N={len(meta)},1536) "
                 f"aligned to {name} row order>.")
    F = np.load(feats_path).astype(np.float32)
    F = F[:, 0] if F.ndim == 3 else F
    assert F.shape == (len(meta), 1536), F.shape
    ok = np.isfinite(F).all(1)
    meta["status"] = np.where(ok, "ok", "no_feature")
    meta["n_context"] = meta.groupby("slide").slide.transform("size").where(ok, 0)
    vocab_genes = pd.read_csv(os.path.join(OUT, "hest_val_genes.csv")).gene.tolist()
    cols = m.gene_cols(vocab_genes)
    pred = np.full((len(meta), len(cols)), np.nan, np.float32)
    for sl, idx in meta[ok].groupby("slide").indices.items():
        idx = np.where(ok)[0][idx]
        pred[idx] = m.predict(F[idx], meta[["x", "y"]].values[idx], meta.organ.values[idx[0]], "Visium", cols)[0]
    np.save(os.path.join(OUT, f"{name}_pred.npy"), pred.astype(np.float16))
    pd.Series(vocab_genes).to_csv(os.path.join(OUT, f"{name}_genes.csv"), index=False, header=["gene"])
    meta.to_csv(os.path.join(OUT, f"{name}_meta.csv"), index_label="row")
    print(f"[{name}] predicted {ok.sum()} / {len(meta)}; {meta.slide.nunique()} slides")
    if name == "orion":
        orion_protein_check(pred, vocab_genes)


ORION_PAIRS = [("PTPRC", "CD45"), ("EPCAM", "Pan-CK"), ("KRT8", "Pan-CK"), ("CDH1", "E-cadherin"), ("ACTA2", "SMA"),
               ("MKI67", "Ki67"), ("PECAM1", "CD31"), ("CD68", "CD68"), ("CD3E", "CD3e"), ("CD8A", "CD8a"),
               ("MS4A1", "CD20"), ("CD163", "CD163")]


def orion_protein_check(pred, genes):
    """predicted RNA vs measured ORION protein (= <marker>_count / n_nuclei) per tile, tiles with >= 20 nuclei."""
    from scipy.stats import pearsonr, spearmanr
    sys.path.insert(0, ROOT)
    from common import load_all_tiles
    d = load_all_tiles()
    ts = pd.read_csv(os.path.join(ROOT, "results", "tile_stats.csv"))
    assert (ts.tile_id.values == d.tile_id.values).all()
    keep = (ts.n_nuclei.values >= 20) & np.isfinite(pred).all(1)
    gi = {g: i for i, g in enumerate(genes)}
    rows = []
    for g, p in ORION_PAIRS:
        x = pred[keep, gi[g]].astype(np.float64)
        y = (d[f"{p}_count"].values / np.maximum(ts.n_nuclei.values, 1))[keep]
        ws = [spearmanr(x[i], y[i])[0] for i in d[keep].groupby("in_slide_name").indices.values() if len(i) >= 20]
        rows.append({"gene": g, "protein": p, "n_tiles": int(keep.sum()), "pearson": pearsonr(x, y)[0],
                     "spearman": spearmanr(x, y)[0], "within_slide_spearman_median": np.nanmedian(ws),
                     "n_slides_within": len(ws)})
    R = pd.DataFrame(rows)
    R.to_csv(os.path.join(OUT, "orion_protein_check.csv"), index=False)
    print(R.round(3).to_string(index=False))


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--set", required=True, choices=["int2_check", "hest_val", "orion", "tcga_text"])
    ap.add_argument("--feats", default=None, help="GigaPath (N,1536) features for orion / tcga_text")
    ap.add_argument("--extract", action="store_true", help="compute GigaPath features for --set orion|tcga_text (needs HF_TOKEN env)")
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    a = ap.parse_args()
    torch.manual_seed(0)
    if a.extract:
        extract_gigapath(a.set, a.device)
        sys.exit(0)
    model = STPath(a.device)
    {"int2_check": lambda: run_int2(model), "hest_val": lambda: run_hest(model)}.get(
        a.set, lambda: run_target(model, a.set, a.feats))()