"""Re-check of the H&E <-> text encoder comparison (20_pairing.py, 28_) under the final protocol with CLEAN real text:
TCGA splits/final/tcga_fold<k>.csv (patient 5-fold, val = early stopping only), junk texts removed (40_clean_text.py),
identical text_head recipe to 33_final_eval.py (frozen image anchor, MLP on frozen text features, multi-positive InfoNCE).
Image anchors x text towers; PathGen-CLIP is left out (trained on PathGen-1.6M = most of this pool -> not held out).
Output results/text_enc_clean_long.csv, text_enc_clean_summary.csv
"""
import os

import numpy as np
import pandas as pd
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

IMGS = ["uni2", "phikonv2", "gigapath", "conch"]
TXTS = ["biomedclip", "plip", "imagebind", "conchclip"]
K, G100, DRAWS = 5, 100, 50
dev = "cuda"
T = pd.read_csv(os.path.join(CACHE, "alltxt_meta.csv"))
t_img = torch.as_tensor(T.img_idx.values, device=dev)
T = T[~pd.read_csv(os.path.join(CACHE, "alltxt_junk.csv")).junk.values]


def mlp(din, dout, hid=512):  # = 33_
    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 one_text_per_image(t_ids, rng):  # = 33_
    pick = T.loc[t_ids].groupby("img_idx").apply(lambda g: g.index[rng.integers(len(g))], include_groups=False)
    return pick.index.values, pick.values


def g100(a, b, seed=0):  # = 33_
    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()}


def train(X, Tx, tr_t, va_t, epochs=60, bs=256, lr=3e-4, seed=0):  # = 33_ train_text_head
    torch.manual_seed(seed); rng = np.random.default_rng(seed)
    th = mlp(Tx.shape[1], X.shape[-1]).to(dev)
    logit = nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32, device=dev))
    opt = torch.optim.AdamW(list(th.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_t) // bs), pct_start=0.05)
    va_i, va_tt = one_text_per_image(va_t, np.random.default_rng(0))
    za_va = F.normalize(X[torch.as_tensor(va_i, device=dev), 0], dim=-1)
    best, state = -1, None
    for ep in range(epochs):
        th.train()
        perm = rng.permutation(tr_t)
        for s in range(0, len(perm) - bs + 1, bs):
            b = torch.as_tensor(perm[s:s + bs], device=dev)
            ii = t_img[b]
            zi = F.normalize(X[ii, torch.randint(0, X.shape[1], (len(b),), device=dev)], dim=-1)
            lg = logit.exp().clamp(max=100) * zi @ F.normalize(th(Tx[b]), dim=-1).T
            pos = (ii[:, None] == ii[None, :]).float()
            tgt = pos / pos.sum(1, keepdim=True)
            loss = (torch.sum(-tgt * F.log_softmax(lg, 1), 1).mean() + torch.sum(-tgt * F.log_softmax(lg.T, 1), 1).mean()) / 2
            opt.zero_grad(set_to_none=True); loss.backward(); opt.step(); sched.step()
        th.eval()
        with torch.no_grad():
            sim = za_va @ F.normalize(th(Tx[torch.as_tensor(va_tt, device=dev)]), dim=-1).T
            v = float((((sim > sim.diag()[:, None]).sum(1) == 0).float().mean() + ((sim.T > sim.diag()[:, None]).sum(1) == 0).float().mean()) / 2)
        if v > best:
            best, state = v, {k_: x.detach().clone() for k_, x in th.state_dict().items()}
    th.load_state_dict(state); th.eval()
    return th, best


rows = []
for ie in IMGS:
    X = torch.from_numpy(np.load(os.path.join(CACHE, f"allimg_he_{ie}.npy")).astype(np.float32)).to(dev)
    for te_ in TXTS:
        Tx = torch.from_numpy(np.load(os.path.join(CACHE, f"alltxt_text_{te_}.npy"))[:, 0].astype(np.float32)).to(dev)
        for k in range(K):
            sp = pd.read_csv(os.path.join(ROOT, "splits", "final", f"tcga_fold{k}.csv")).set_index("img_idx").split
            tsp = sp.loc[T.img_idx].values
            tr_t, va_t, te_t = (T.index[tsp == x].values for x in ["train", "val", "test"])
            th, best = train(X, Tx, tr_t, va_t)
            ii, tt = one_text_per_image(te_t, np.random.default_rng(k))
            with torch.no_grad():
                za = F.normalize(X[torch.as_tensor(ii, device=dev), 0], dim=-1).cpu()
                zt = F.normalize(th(Tx[torch.as_tensor(tt, device=dev)]), dim=-1).cpu()
            rows.append({"img_enc": ie, "txt_enc": te_, "fold": k, "val_R@1": best, "n_test": len(ii), **g100(za, zt)})
            print(f"{ie} + {te_}  fold {k}: val {best:.3f}  test R@1 {rows[-1]['R@1']:.3f}", flush=True)
R = pd.DataFrame(rows)
R.to_csv(os.path.join(ROOT, "results", "text_enc_clean_long.csv"), index=False)
S = R.groupby(["img_enc", "txt_enc"], sort=False)[["R@1", "R@10", "pair_acc"]].agg(["mean", "std"]).round(4)
S.to_csv(os.path.join(ROOT, "results", "text_enc_clean_summary.csv"))
pd.set_option("display.width", 200)
print(R.pivot_table(index="img_enc", columns="txt_enc", values="R@1", sort=False).round(3).to_string())