"""Evaluate a saved checkpoint from the cached encoder features through the public embed.py API (no frozen encoders,
no raw data; needs cache/ + splits/ from the features repo).

    python -I evaluate_saved.py [models/ej_ad_mv_fold0]

Uses the held-out fold-0 test items stored with the checkpoint (test_items.npz, written by 52_final_model.py fold0)
and the same g100 protocol as 52_ (100-item galleries, 50 draws, seed 0, both directions; Hungarian pair accuracy).
If the checkpoint has a test.csv, the numbers must match it (pure inference: deterministic up to float rounding).
Output: <model>/test_recomputed.csv
"""
import os
import sys

import numpy as np
import pandas as pd
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")
CACHE = os.path.join(ROOT, "cache")
H, RF = os.path.join(CACHE, "hest"), os.path.join(CACHE, "rnafm")
ld = lambda *p: np.load(os.path.join(*p), mmap_mode="r")
E = Embedder(MD)
TI = np.load(os.path.join(MD, "test_items.npz"))
o, hs, ii, tt = TI["orion"], TI["hest"], TI["tcga_img"], TI["tcga_txt"]


def g100(a, b, seed=0, G=100, draws=50):  # = 52_final_model.g100
    from scipy.optimize import linear_sum_assignment
    rng = np.random.default_rng(seed)
    out = {"R@1": [], "R@10": [], "pair_acc": []}
    for _ in range(draws):
        s = torch.from_numpy(rng.choice(len(a), min(G, 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)
            out["R@1"].append((rank < 1).mean() / 2); out["R@10"].append((rank < 10).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()}


def virt_hvg(st):  # STPath virtual RNA -> log1p CP10k on the model's HVG list (= 52_)
    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 c[:, pd.Series(range(len(g)), index=g).loc[E.hvg].values]


def rna(st, idx):
    if st == "hest":
        F = {e: (ld(H, "hvg.npy") if e == "hvg" else ld(RF, f"{e}_hest.npy"))[idx] for e in RNA_E}
    else:
        F = {e: (virt_hvg({"orion_v": "orion", "tcga_v": "tcga_text"}[st]) if e == "hvg" else ld(RF, f"{e}_{st}.npy"))[idx]
             for e in RNA_E}
    return E.rna_from_feats(F)


text = lambda st, idx: E.text_from_feats({e: ld(CACHE, f"{st}_text_{e}.npy")[idx, 0] for e in TXT_E})
zm = E.mif_from_feats(ld(CACHE, "if_virtues_m16.npy")[o, 0])
zt = text("alltxt", tt)
res = {"H&E<->mIF (ORION)": g100(E.he_from_feats(ld(CACHE, "he_uni2.npy")[o, 0], ld(CACHE, "he_medvtok.npy")[o, 0]), zm),
       "H&E<->RNA (HEST)": g100(E.he_from_feats(ld(H, "he_uni2.npy")[hs, 0], ld(H, "he_medvtok.npy")[hs, 0]), rna("hest", hs)),
       "H&E<->text (TCGA)": g100(E.he_from_feats(ld(CACHE, "allimg_he_uni2.npy")[ii, 0],
                                                 ld(CACHE, "allimg_he_medvtok.npy")[ii, 0]), zt),
       "mIF<->RNA(virtual) (ORION)": g100(zm, rna("orion_v", o)),
       "mIF<->text(synthetic) (ORION)": g100(zm, text("orion_gentext", o)),
       "text<->RNA(virtual) (TCGA)": g100(zt, rna("tcga_v", ii)),
       "text<->mIF(MIPHEI) (TCGA)": g100(zt, E.mif_from_feats(ld(CACHE, "if_virtues_m16_virtual_miphei_tcga_text.npy")[ii, 0]))}
D = pd.DataFrame([{"pair": p, **r} for p, r in res.items()])
ref = os.path.join(MD, "test.csv")
if os.path.exists(ref):
    D["R@1_saved"] = pd.read_csv(ref).set_index("pair").loc[D.pair, "R@1"].values
    D["abs_diff"] = (D["R@1"] - D["R@1_saved"]).abs()
D.to_csv(os.path.join(MD, "test_recomputed.csv"), index=False)
pd.set_option("display.width", 200)
print(f"{MD}: fold-0 test, g100 (chance R@1 0.01)\n{D.round(4).to_string(index=False)}")
if "abs_diff" in D:
    ok = D.abs_diff.max() < 2e-3
    print("MATCHES the saved test.csv" if ok else "DOES NOT MATCH the saved test.csv", f"(max |diff| {D.abs_diff.max():.2e})")
    sys.exit(0 if ok else 1)