"""H&E foundation-model ablation (UNI2 vs GigaPath vs CONCH as the frozen ImageBind anchor).

  --prep : cache/pathgen_he_<anchor>.npy from cache/allimg_he_<anchor>.npy (PathGen-1.6M captions are a subset of the
           all-pairs pool, same image files -> identical features; verified for uni2, cosine 1.0)
  default: evaluate
    (A) ORION H&E<->mIF, trained pair: runs/bench_fm/*_<anchor>_<if_enc>_fm-geneformer_pathgenclip (17_bench_fm.py);
        tie-aware R@1 on 100-tile galleries + Hungarian pairing acc on the full test fold, both directions
    (B) emergent text<->mIF: text head trained ONLY on PathGen H&E-caption pairs with the same anchor
        (runs/text_pathgen/<anchor>_<txt>, 12_train_text_pathgen.py) scores the concept prompts against the ORION
        mIF embeddings of that anchor's run (text and mIF never co-occur in training); AUROC + Spearman (18_eval_fm.probe)
Outputs results/anchor_ifret.csv, results/anchor_text_probes.csv
"""
import glob
import json
import os
import sys

import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F

from common import CACHE, ROOT, RUNS, metrics

ANCHORS = ["uni2", "gigapath", "conch"]
OUT = os.path.join(ROOT, "results")

if "--prep" in sys.argv:
    P = pd.read_csv(os.path.join(CACHE, "pathgen_meta.csv"))
    I = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv")).set_index("key")
    ix = I.loc[P.image_path.map(lambda p: os.path.basename(p).rsplit(".", 1)[0]), "idx"].values
    assert (I.image_path.values[ix] == P.image_path.values).all()
    for a in ANCHORS:
        f = os.path.join(CACHE, f"pathgen_he_{a}.npy")
        if not os.path.exists(f):
            np.save(f, np.load(os.path.join(CACHE, f"allimg_he_{a}.npy"))[ix])
            print("saved", f)
    sys.exit()

# reuse the probe of 18_eval_fm.py without running its main block
src = open(os.path.join(ROOT, "18_eval_fm.py")).read().split("# collect test embeddings")[0]
ns = {}
exec(compile(src, "18_eval_fm_head", "exec"), ns)
probe, prompts = ns["probe"], ns["prompts"]


def load_head(anchor, enc, din, dout):
    h = torch.nn.Sequential(torch.nn.Linear(din, 512), torch.nn.GELU(), torch.nn.Dropout(0.1), torch.nn.Linear(512, 512),
                            torch.nn.GELU(), torch.nn.LayerNorm(512), torch.nn.Linear(512, dout))
    h.load_state_dict(torch.load(os.path.join(RUNS, "text_pathgen", f"{anchor}_{enc}", "head.pt"), map_location="cpu"))
    return h.eval()


ret, prb = [], []
for f in sorted(glob.glob(os.path.join(RUNS, "bench_fm", "*", "test_embeddings.pt"))):
    a = json.load(open(os.path.join(os.path.dirname(f), "train.json")))["args"]
    if a["anchor"] not in ANCHORS or a["mk_enc"] != "fm:geneformer" or a["txt_enc"] != "pathgenclip":
        continue
    E = torch.load(f, weights_only=False)  # our own files (contain numpy tile ids)
    he, zi, tiles = E["z"]["he"], E["z"]["if"], np.asarray(E["tile_id"])
    key = {"anchor": a["anchor"], "if_enc": a["if_enc"], "split": a["split"], "fold": a["fold"]}
    for q, g, d in [(he, zi, "he->if"), (zi, he, "if->he")]:
        ret.append({**key, "dir": d, **metrics(q, g)})
    for t in ["biomedclip", "plip", "pathgenclip"]:
        if not os.path.exists(os.path.join(RUNS, "text_pathgen", f"{a['anchor']}_{t}", "head.pt")):
            continue
        prm = torch.from_numpy(np.load(os.path.join(CACHE, f"prompts_text_{t}.npy"))[:, 0].astype(np.float32))
        with torch.no_grad():
            zp = F.normalize(load_head(a["anchor"], t, prm.shape[1], he.shape[1])(prm), dim=-1)
        for m, z in [("he", he), ("if", zi)]:
            for i, (c, pr) in enumerate(prompts.items()):
                prb.append({**key, "text_enc": t, "modality": m, "concept": c, **probe((z @ zp[i]).numpy(), tiles, c, pr["label"])})

R, P = pd.DataFrame(ret), pd.DataFrame(prb)
R.to_csv(os.path.join(OUT, "anchor_ifret.csv"), index=False)
P.to_csv(os.path.join(OUT, "anchor_text_probes.csv"), index=False)
pd.set_option("display.width", 250)
g = R.groupby(["if_enc", "anchor", "split", "fold"])[["R@1_g100", "pair_acc"]].mean().reset_index()
print(g.groupby(["if_enc", "anchor", "split"])[["R@1_g100", "pair_acc"]].agg(["mean", "std"]).round(3).to_string())
if len(P):
    print(P.groupby(["text_enc", "modality", "if_enc", "anchor"])[["auroc", "spearman"]].mean().round(3).to_string())