"""End-to-end check of embed.py from RAW inputs on examples/ (64 held-out fold-0 test items per dataset; needs the
frozen encoders, see setup_encoders.sh; GPU).

    python -I example_check.py [models/ej_ad_mv_fold0]

1. raw input -> encoder features, compared with the cached features the model was trained / tested on
   (examples/ref_feats.npz): per-item cosine. Expected ~1.000; Tahoe-x1 / SE-600M ~0.99 (their official input
   pipelines sample genes at random); VirTues lower here because mIF standardisation statistics are computed over the
   tiles given in one call (64 tiles here vs the whole ORION set for the cache).
2. raw input -> 1536-d embeddings -> retrieval inside the 64 items (R@1, both directions; chance 1/64), next to the
   same retrieval from the cached features.
Output: examples/example_check.csv; exit code 1 if a deterministic encoder does not reproduce its cached features.
"""
import os
import sys
import time

import numpy as np
import pandas as pd
import scipy.sparse as sp
import torch

ROOT = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, ROOT)
from embed import RNA_E, TXT_E, Embedder  # noqa: E402

MD = sys.argv[1] if len(sys.argv) > 1 else os.path.join(ROOT, "models", "ej_ad_mv_fold0")
EX = os.path.join(ROOT, "examples")
E = Embedder(MD)
R = dict(np.load(os.path.join(EX, "ref_feats.npz")))
rows, t0 = [], time.time()
MIN_COS = {"uni2": .999, "medvtok": .999, "keep": .999, "biomedclip": .99, "qwen3-0.6b": .999, "hvg": .999,
           "aidocell": .999, "tahoex1": .97, "se600m": .97, "virtues": 0.0}  # median cosine required (virtues: report only)


def cos(a, b):
    a, b = torch.as_tensor(np.asarray(a, np.float32)), torch.as_tensor(np.asarray(b, np.float32))
    return torch.nn.functional.cosine_similarity(a, b, dim=-1).numpy()


def feat(ds, enc, raw):
    c = cos(raw, R[f"{ds}_{enc}"])
    rows.append({"check": "features", "data": ds, "encoder": enc, "n": len(c), "cos_median": float(np.median(c)),
                 "cos_min": float(c.min()), "ok": bool(np.median(c) >= MIN_COS[enc])})
    print(f"[{time.time() - t0:4.0f}s] {ds:5s} {enc:11s} cos median {np.median(c):.4f} min {c.min():.4f}"
          f"{'' if rows[-1]['ok'] else '   <-- FAIL'}", flush=True)


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


# raw H&E (PNG files / uint8 arrays) -> UNI2 + MedVTok
U, M = {}, {}
U["orion"], M["orion"] = E.he_feats(sorted(os.path.join(EX, "orion_he", f) for f in os.listdir(os.path.join(EX, "orion_he"))))
U["hest"], M["hest"] = E.he_feats(np.load(os.path.join(EX, "hest_he.npy")))
C = pd.read_csv(os.path.join(EX, "tcga_captions.csv"))
U["tcga"], M["tcga"] = E.he_feats([os.path.join(EX, "tcga_he", f) for f in C.file])
for ds in U:
    feat(ds, "uni2", U[ds]); feat(ds, "medvtok", M[ds])
E.release()
# raw text -> KEEP / BiomedCLIP / Qwen3
Ft = E.text_feats(C.caption.fillna("").astype(str).tolist())
for e in TXT_E:
    feat("tcga", e, Ft[e])
E.release()
# raw counts -> hvg / Tahoe-x1 / SE-600M / AIDO.Cell
genes = pd.read_csv(os.path.join(EX, "hest_genes.csv")).gene.values
Fr = E.rna_feats(sp.load_npz(os.path.join(EX, "hest_counts.npz")), genes)
for e in RNA_E:
    feat("hest", e, Fr[e])
E.release()
# raw mIF tiles -> VirTues
m = np.load(os.path.join(EX, "orion_mif.npz"), allow_pickle=True)
Fm = E.mif_feats(m["tiles"], m["slide"])
feat("orion", "virtues", Fm)
E.release()

# embeddings + retrieval inside the 64 items
pairs = {"H&E<->mIF (ORION)": ((E.he_from_feats(U["orion"], M["orion"]), E.mif_from_feats(Fm)),
                               (E.he_from_feats(R["orion_uni2"], R["orion_medvtok"]), E.mif_from_feats(R["orion_virtues"]))),
         "H&E<->RNA (HEST)": ((E.he_from_feats(U["hest"], M["hest"]), E.rna_from_feats(Fr)),
                              (E.he_from_feats(R["hest_uni2"], R["hest_medvtok"]),
                               E.rna_from_feats({e: R[f"hest_{e}"] for e in RNA_E}))),
         "H&E<->text (TCGA)": ((E.he_from_feats(U["tcga"], M["tcga"]), E.text_from_feats(Ft)),
                               (E.he_from_feats(R["tcga_uni2"], R["tcga_medvtok"]),
                                E.text_from_feats({e: R[f"tcga_{e}"] for e in TXT_E})))}
for p, ((a, b), (ac, bc)) in pairs.items():
    rows.append({"check": "retrieval", "data": p, "n": len(a), "R@1_raw": r1(a, b), "R@1_cached": r1(ac, bc),
                 "chance": 1 / len(a)})
    print(f"{p:18s} R@1 within {len(a)}: raw inputs {rows[-1]['R@1_raw']:.3f} | cached features "
          f"{rows[-1]['R@1_cached']:.3f} (chance {1 / len(a):.3f})", flush=True)
D = pd.DataFrame(rows)
D.to_csv(os.path.join(EX, "example_check.csv"), index=False)
bad = D[(D.check == "features") & ~D.ok.astype(bool)]
print(f"done [{time.time() - t0:.0f}s]:", "all encoders reproduce the cached features" if bad.empty else
      f"FAILED {bad[['data', 'encoder']].values.tolist()}")
sys.exit(0 if bad.empty else 1)