"""External test sets (never used for training): DLPFC Visium (brain, 12 sections / 3 donors) and HEST IDC
(breast, Xenium pseudo-spots, 4 samples; 280-gene panel). Prepared by 36_external_prep.py -> cache/external/.

Real data only for training (HEST-1k, train/val by study|patient group, val = early stopping only):
  dlpfc head : H&E(UNI2) <-> RNA MLP on the 2000 HEST HVGs, trained on HEST WITHOUT the spatialLIBD study
               (= DLPFC; those spots are in cache/hest) -> DLPFC is fully held out
  idc head   : same recipe on the IDC panel genes present in the HEST common gene set (IDC samples are not in HEST cache)
  RNA input  : log1p(CP10k over the HEST common genes / over the panel genes), then z-scored per gene with the
               statistics of its OWN dataset (HEST training spots vs the external set; label-free platform alignment)
  5 heads per set (different val groups / seeds) -> mean +- std.
Virtual data only at test time: MIPHEI-ViT virtual mIF of the external H&E (35_ + 30_ --m16) through the 16-ch mIF
heads trained on REAL ORION (patient 5-fold, fold k).
Synthetic text only at test time: GPT captions of the external H&E (27_ --sets dlpfc idc, 32_, 08_ biomedclip)
through the TCGA text heads (runs/text_clean/uni2_biomedclip_fold<k>, trained on CLEAN real image-text pairs in 33_).
Pairs: H&E<->RNA (real, external), RNA<->mIF(MIPHEI) (emergent, real x virtual), H&E<->mIF(MIPHEI) (circular),
       RNA<->text(synthetic) (emergent, real x virtual), H&E<->text(synthetic) (circular),
       mIF(MIPHEI)<->text(synthetic) (virtual x virtual).
Metrics: 100-item galleries (50 draws, both directions): R@1/R@5/R@10/MRR + Hungarian pairing accuracy.
Usage: python 38_external_eval.py feats   (UNI2 features of the external patches, GPU)
       python 38_external_eval.py eval    -> results/external_long.csv, external_summary.csv
"""
import os
import sys

import numpy as np
import pandas as pd
import scipy.sparse as sp
import torch
import torch.nn as nn
import torch.nn.functional as F
from scipy.optimize import linear_sum_assignment

from common import CACHE, ROOT, RUNS

E = os.path.join(CACHE, "external")
H = os.path.join(CACHE, "hest")
OUT = os.path.join(ROOT, "results")
SETS = {"dlpfc": "dlpfc_he_05um.npy"}  # IDC dropped from the report 2026-10-18 (user request)
MODE = sys.argv[1]

if MODE == "feats":
    import timm
    dev = "cuda"
    model = timm.create_model(  # = 04b_extract_anchor_fm.py uni2 (local official UNI2-h checkpoint)
        "vit_giant_patch14_224", pretrained=False, img_size=224, patch_size=14, depth=24, num_heads=24,
        init_values=1e-5, embed_dim=1536, mlp_ratio=2.66667 * 2, num_classes=0, no_embed_class=True,
        mlp_layer=timm.layers.SwiGLUPacked, act_layer=torch.nn.SiLU, reg_tokens=8, dynamic_img_size=True)
    model.load_state_dict(torch.load("/diskarray/home/tl688/reg2026/model_alldata/uni2-h.bin", map_location="cpu"), strict=True)
    model = model.eval().to(dev)
    mean = torch.tensor([0.485, 0.456, 0.406], device=dev).view(1, 3, 1, 1)
    std = torch.tensor([0.229, 0.224, 0.225], device=dev).view(1, 3, 1, 1)
    for st, f in SETS.items():
        X = np.load(os.path.join(E, f), mmap_mode="r")
        out = np.zeros((len(X), 1536), np.float32)
        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16):
            for s in range(0, len(X), 32):
                x = torch.from_numpy(np.asarray(X[s:s + 32])).to(dev).permute(0, 3, 1, 2).float() / 255
                if x.shape[-1] != 224:
                    x = F.interpolate(x, size=(224, 224), mode="bilinear", antialias=True, align_corners=False)
                out[s:s + 32] = model((x - mean) / std).float().cpu().numpy()
        np.save(os.path.join(E, f"{st}_he_uni2.npy"), out)
        print(st, out.shape, flush=True)
    # protocol check: DLPFC spots that are also in the HEST cache (same barcode) -> cosine to the HEST UNI2 feature
    m, hm = pd.read_csv(os.path.join(E, "dlpfc_meta.csv")), pd.read_csv(os.path.join(H, "meta.csv"))
    hf = np.load(os.path.join(H, "he_uni2.npy"))[:, 0].astype(np.float32)
    o = np.load(os.path.join(E, "dlpfc_he_uni2.npy"))
    j = m.reset_index().merge(hm.reset_index(), left_on=["hest_id", "barcode"], right_on=["sample", "spot"], suffixes=("_e", "_h"))
    if len(j):
        a, b = F.normalize(torch.from_numpy(o[j.index_e.values]), dim=-1), F.normalize(torch.from_numpy(hf[j.index_h.values]), dim=-1)
        rnd = F.normalize(torch.from_numpy(hf[np.random.default_rng(0).choice(len(hf), 200, replace=False)]), dim=-1)
        print(f"protocol check: {len(j)} shared spots, cosine(ours, HEST cache) = {(a * b).sum(1).numpy().round(3)}; "
              f"mean cosine to random HEST spots {float((a @ rnd.T).mean()):.3f}")
    sys.exit(0)

# ======================================================================== eval
dev = "cuda" if torch.cuda.is_available() else "cpu"
G100, DRAWS, K = 100, 50, 5
HM = pd.read_csv(os.path.join(H, "meta.csv"))
hgenes = pd.read_csv(os.path.join(H, "genes.csv")).gene.values
hcounts = sp.load_npz(os.path.join(H, "counts.npz")).tocsr().astype(np.float32)
hvg = pd.read_csv(os.path.join(H, "hvg_genes.csv")).gene.values
HE_H = torch.from_numpy(np.load(os.path.join(H, "he_uni2.npy"))[:, 0].astype(np.float32))


def mlp(din, dout=1536, hid=512):
    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 lognorm(C, total):
    return np.log1p(C / np.maximum(total, 1)[:, None] * 1e4).astype(np.float32)


def zscore(x, ref):
    mu, sd = ref.mean(0), ref.std(0) + 1e-6
    return (x - mu) / sd


def r1(a, b):
    s = a @ b.T
    return float(((s > s.diag()[:, None]).sum(1) == 0).float().mean())


def train_head(G, tr, va, seed, epochs=100, bs=256, lr=3e-4):  # = 23_bind_gene.py recipe
    torch.manual_seed(seed); rng = np.random.default_rng(seed)
    A, Gd = HE_H.to(dev), torch.from_numpy(G).to(dev)
    head = mlp(G.shape[1]).to(dev)
    logit = nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32, device=dev))
    opt = torch.optim.AdamW(list(head.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) // bs), pct_start=0.05)
    za_va = F.normalize(A[va], dim=-1)
    best, state = -1, None
    for ep in range(epochs):
        head.train()
        perm = rng.permutation(tr)
        for s in range(0, len(perm) - bs + 1, bs):
            b = torch.as_tensor(perm[s:s + bs], device=dev)
            lg = logit.exp().clamp(max=100) * F.normalize(A[b], dim=-1) @ F.normalize(head(Gd[b]), dim=-1).T
            t = torch.arange(len(b), device=dev)
            loss = (F.cross_entropy(lg, t) + F.cross_entropy(lg.T, t)) / 2
            opt.zero_grad(set_to_none=True); loss.backward(); opt.step(); sched.step()
        head.eval()
        with torch.no_grad():
            zg = F.normalize(head(Gd[va]), dim=-1)
            v = (r1(za_va, zg) + r1(zg, za_va)) / 2
        if v > best:
            best, state = v, {k_: x.detach().cpu().clone() for k_, x in head.state_dict().items()}
    head = head.cpu(); head.load_state_dict(state); head.eval()
    return head, best


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)
            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()}


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


def mif_head(k):
    d = torch.load(os.path.join(RUNS, "bench_fm", f"patient_5fold_fold{k}_uni2_fm-virtues_m16_fm-geneformer_pathgenclip", "heads.pt"), map_location="cpu")
    h = mlp(d["state"]["if"]["0.weight"].shape[1]); h.load_state_dict(d["state"]["if"]); h.eval()
    return lambda x: h((torch.as_tensor(x, dtype=torch.float32) - d["if_mean"]) / d["if_std"])


def text_head(k):
    sd = torch.load(os.path.join(RUNS, "text_clean", f"uni2_biomedclip_fold{k}", "head.pt"), map_location="cpu")
    h = mlp(sd["0.weight"].shape[1], sd["6.weight"].shape[0]); h.load_state_dict(sd); h.eval()
    return lambda x: h(torch.as_tensor(x, dtype=torch.float32))


rows = []
hset = pd.Series(range(len(hgenes)), index=hgenes)
htot = np.asarray(hcounts.sum(1)).ravel()
for st in SETS:
    em = pd.read_csv(os.path.join(E, f"{st}_meta.csv"))
    he = np.load(os.path.join(E, f"{st}_he_uni2.npy"))
    if st == "dlpfc":
        eg = pd.read_csv(os.path.join(E, "dlpfc_genes_full.csv")).gene.values
        ec = np.load(os.path.join(E, "dlpfc_counts_full.npy")).astype(np.float32)
        eidx = pd.Series(range(len(eg)), index=eg)
        eidx = eidx[~eidx.index.duplicated()]
        common = [g for g in hgenes if g in eidx.index]
        genes = [g for g in hvg if g in eidx.index]
        etot = ec[:, eidx.loc[common].values].sum(1)                      # library size over the HEST common genes
        excl = HM.group.str.startswith("spatialLIBD").values                # DLPFC study removed from training
        htot_ = htot
    else:
        eg = pd.read_csv(os.path.join(E, "idc_genes.csv")).gene.values
        ec = np.load(os.path.join(E, "idc_counts.npy")).astype(np.float32)
        eidx = pd.Series(range(len(eg)), index=eg)
        eidx = eidx[~eidx.index.duplicated()]
        genes = [g for g in eg if g in hset.index]
        etot = ec[:, eidx.loc[genes].values].sum(1)                         # library size over the panel genes
        htot_ = np.asarray(hcounts[:, hset.loc[genes].values].sum(1)).ravel()
        excl = HM["sample"].isin(em["sample"].unique()).values
    Gh = lognorm(hcounts[:, hset.loc[genes].values].toarray(), htot_)
    Ge = lognorm(ec[:, eidx.loc[genes].values], etot)
    pool = np.flatnonzero(~excl & (htot_ > 0))
    print(f"{st}: {len(em)} external spots, {len(genes)} genes, HEST training pool {len(pool)} spots "
          f"({int(excl.sum())} excluded)", flush=True)
    f_mv = os.path.join(CACHE, f"if_virtues_m16_virtual_miphei_{st}.npy")
    mv = np.load(f_mv)[:, 0].astype(np.float32) if os.path.exists(f_mv) else None
    f_t = os.path.join(CACHE, f"{st}_gentext_text_biomedclip.npy")
    tx = np.load(f_t)[:, 0].astype(np.float32) if os.path.exists(f_t) else None
    zh = nz(he)
    for k in range(K):
        grp = np.random.default_rng(k).permutation(sorted(HM.group.iloc[pool].unique()))
        va_g = set(grp[: max(1, len(grp) // 10)])
        isva = HM.group.iloc[pool].isin(va_g).values
        tr, va = pool[~isva], pool[isva]
        for norm in ["z-scored per dataset", "plain"]:
            Gtr = zscore(Gh, Gh[tr]) if norm != "plain" else Gh
            Gex = zscore(Ge, Ge) if norm != "plain" else Ge
            head, best = train_head(Gtr, tr, va, seed=k)
            with torch.no_grad():
                zg = nz(head(torch.from_numpy(Gex)))
            rows.append({"set": st, "fold": k, "rna_input": norm, "part": "real", "pair": "H&E<->RNA",
                         "kind": "trained pair (external)", "circular": False, "n": len(zh), "val_R@1": best, **g100(zh, zg)})
            if mv is not None and norm == "z-scored per dataset":
                with torch.no_grad():
                    zm = nz(mif_head(k)(mv))
                for pair, a, b, kind, circ in [("RNA<->mIF(MIPHEI)", zg, zm, "emergent: real x virtual", False),
                                               ("H&E<->mIF(MIPHEI)", zh, zm, "virtual from this H&E", True)]:
                    rows.append({"set": st, "fold": k, "rna_input": norm, "part": "virtual", "pair": pair, "kind": kind,
                                 "circular": circ, "n": len(zh), **g100(a, b)})
            if tx is not None and norm == "z-scored per dataset":
                with torch.no_grad():
                    zt = nz(text_head(k)(tx))
                    zm = nz(mif_head(k)(mv)) if mv is not None else None
                for pair, a, b, kind, circ in [("RNA<->text(synthetic)", zg, zt, "emergent: real x virtual", False),
                                               ("H&E<->text(synthetic)", zh, zt, "virtual from this H&E", True),
                                               ("mIF(MIPHEI)<->text(synthetic)", zm, zt, "emergent: virtual x virtual", True)]:
                    if a is not None:
                        rows.append({"set": st, "fold": k, "rna_input": norm, "part": "virtual", "pair": pair, "kind": kind,
                                     "circular": circ, "n": len(zh), **g100(a, b)})
            if norm == "z-scored per dataset":  # embeddings for 39_external_check.py
                np.savez(os.path.join(E, f"emb_{st}_fold{k}.npz"), he=zh.numpy(), rna=zg.numpy(),
                         **({"mif": zm.numpy()} if mv is not None else {}), **({"text": zt.numpy()} if tx is not None else {}))
        print(f"  {st} head {k}: val R@1 {best:.3f}", flush=True)

R = pd.DataFrame(rows)
R.to_csv(os.path.join(OUT, "external_long.csv"), index=False)
cols = ["R@1", "R@5", "R@10", "MRR", "pair_acc"]
S = R.groupby(["set", "part", "pair", "rna_input", "kind", "circular"], sort=False)[cols + ["n"]].agg(["mean", "std"])
S.round(4).to_csv(os.path.join(OUT, "external_summary.csv"))
pd.set_option("display.width", 250)
print(R.groupby(["set", "part", "pair", "rna_input", "circular"], sort=False)[cols + ["n"]].mean().round(3).to_string())