"""Emergent text alignment on ORION. The text head was trained ONLY on PathGen H&E-caption pairs (12_), the IF /
marker encoders ONLY on ORION H&E pairs (09_, --if_enc convnext_tiny --mk_enc mlp). Text and IF/markers never
co-occur in training; any text<->IF / text<->markers agreement is emergent through the shared UNI2 H&E space.

(A) concept probes: prompts.json (PathGen-style sentences) scored against every ORION test tile embedding
    (he / if / markers); ground truth = marker-positive fraction (or nuclei count). Metrics: Spearman rho and
    AUROC (top quartile vs rest), on tissue tiles (>=20 nuclei; all tiles for 'cellularity').
    Baselines: zero-shot PLIP and ImageBind (their own image/text towers, H&E only, no training).
(B) retrieval with templated ORION descriptions (07_) as text queries through the PathGen-trained head
    (domain-shifted: template style != PathGen captions), tie-aware R@1 on 100-tile galleries.
Outputs results/text_emergent_probes.csv, results/text_emergent_retrieval.csv
"""
import glob
import json
import os

import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
from scipy.stats import spearmanr
from sklearn.metrics import roc_auc_score

from common import CACHE, ROOT, RUNS, load_all_tiles, metrics

TXT = ["plip", "biomedclip", "imagebind"]
OUT = os.path.join(ROOT, "results")
prompts = json.load(open(os.path.join(ROOT, "prompts.json")))
df = load_all_tiles()
stats = pd.read_csv(os.path.join(OUT, "tile_stats.csv")).set_index("tile_id")
df["n_nuclei"] = stats.loc[df.tile_id, "n_nuclei"].values


def load_head(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"uni2_{enc}", "head.pt"), map_location="cpu"))
    return h.eval()


def probe(scores, tiles, key, label_col):
    sub = tiles if key == "cellularity" else tiles[df.loc[tiles, "n_nuclei"].values >= 20]
    pos = np.isin(tiles, sub)
    s = scores[pos]
    y = df.loc[sub, label_col].values.astype(float)
    if label_col != "n_nuclei":
        y = y / np.maximum(df.loc[sub, "n_nuclei"].values, 1)
    lab = y >= np.quantile(y, 0.75)
    return {"spearman": float(spearmanr(s, y).correlation),
            "auroc": float(roc_auc_score(lab, s)) if 0 < lab.sum() < len(lab) else np.nan, "n": int(len(sub))}


probe_rows, ret_rows = [], []
for f in sorted(glob.glob(os.path.join(RUNS, "bench", "*_uni2_convnext_tiny_mlp_*", "test_embeddings.pt"))):
    run = os.path.basename(os.path.dirname(f))
    split, fold = run.split("_fold")[0], int(run.split("_fold")[1].split("_")[0])
    E = torch.load(f, weights_only=False)
    tiles, Z = np.asarray(E["tile_id"]), E["z"]
    for enc in TXT:
        prm = torch.from_numpy(np.load(os.path.join(CACHE, f"prompts_text_{enc}.npy"))[:, 0].astype(np.float32))
        orion_txt = torch.from_numpy(np.load(os.path.join(CACHE, f"text_{enc}.npy"))[tiles, 0].astype(np.float32))
        head = load_head(enc, prm.shape[1], Z["he"].shape[1])
        with torch.no_grad():
            zp = F.normalize(head(prm), dim=-1)
            zt = F.normalize(head(orion_txt), dim=-1)
        for i, (key, pr) in enumerate(prompts.items()):
            for m in ["he", "if", "markers"]:
                sc = (Z[m] @ zp[i]).numpy()
                probe_rows.append({"split": split, "fold": fold, "text_enc": enc, "concept": key, "modality": m,
                                   "kind": "bound (PathGen text head)", **probe(sc, tiles, key, pr["label"])})
        for m in ["he", "if", "markers"]:
            for q, g, name in [(zt, Z[m], f"text->{m}"), (Z[m], zt, f"{m}->text")]:
                ret_rows.append({"split": split, "fold": fold, "text_enc": enc, "pair": name, **metrics(q, g)})
    # zero-shot baselines on H&E with the model's own towers (computed once per fold)
    for fm in ["plip", "imagebind"]:
        hz = F.normalize(torch.from_numpy(np.load(os.path.join(CACHE, f"he_{fm}.npy"))[tiles, 0].astype(np.float32)), dim=-1)
        pz = F.normalize(torch.from_numpy(np.load(os.path.join(CACHE, f"prompts_text_{fm}.npy"))[:, 0].astype(np.float32)), dim=-1)
        for i, (key, pr) in enumerate(prompts.items()):
            probe_rows.append({"split": split, "fold": fold, "text_enc": fm, "concept": key, "modality": "he",
                               "kind": "zero-shot (own image tower)", **probe((hz @ pz[i]).numpy(), tiles, key, pr["label"])})

P = pd.DataFrame(probe_rows); R = pd.DataFrame(ret_rows)
P.to_csv(os.path.join(OUT, "text_emergent_probes.csv"), index=False)
R.to_csv(os.path.join(OUT, "text_emergent_retrieval.csv"), index=False)
pd.set_option("display.width", 250); pd.set_option("display.max_rows", 300)
print(P.groupby(["kind", "text_enc", "modality", "split"])[["auroc", "spearman"]].mean().round(3).unstack("split").to_string())
print(P.groupby(["kind", "text_enc", "modality", "concept"])["auroc"].mean().unstack("concept").round(2).to_string())
print(R.groupby(["text_enc", "pair", "split"])["R@1_g100"].mean().unstack("split").round(3).to_string())