"""Bind real pathology text to the frozen H&E anchor using PathGen-1.6M H&E-caption pairs (ImageBind recipe:
text is paired only with images; it never sees ORION, IF or markers).

frozen text tower features (cache/pathgen_text_<enc>.npy) -> trainable MLP head -> anchor space
loss: symmetric InfoNCE vs frozen anchor image features (cache/pathgen_he_<anchor>.npy, 8 dihedral views)
Usage: python 12_train_text_pathgen.py --txt_enc plip|biomedclip|imagebind [--anchor uni2]
Saves runs/text_pathgen/<anchor>_<enc>/{head.pt, metrics.json}
"""
import argparse
import json
import os

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

from common import CACHE, RUNS

p = argparse.ArgumentParser()
p.add_argument("--txt_enc", default="plip")
p.add_argument("--anchor", default="uni2")
p.add_argument("--epochs", type=int, default=100)
p.add_argument("--bs", type=int, default=256)
p.add_argument("--lr", type=float, default=3e-4)
p.add_argument("--seed", type=int, default=0)
args = p.parse_args()
torch.manual_seed(args.seed); np.random.seed(args.seed)
dev = "cuda"
out_dir = os.path.join(RUNS, "text_pathgen", f"{args.anchor}_{args.txt_enc}")
os.makedirs(out_dir, exist_ok=True)

meta = pd.read_csv(os.path.join(CACHE, "pathgen_meta.csv"))
HE = torch.from_numpy(np.load(os.path.join(CACHE, f"pathgen_he_{args.anchor}.npy")).astype(np.float32))
TX = torch.from_numpy(np.load(os.path.join(CACHE, f"pathgen_text_{args.txt_enc}.npy"))[:, 0].astype(np.float32))
D = HE.shape[-1]
idx = {s: meta.idx[meta.split == s].values for s in ["train", "val", "test"]}


def make_head(din, dout, hid=512):  # same architecture as the ORION benchmark text head
    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))


head = make_head(TX.shape[1], D).to(dev)
logit = nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32, device=dev))
opt = torch.optim.AdamW(list(head.parameters()) + [logit], lr=args.lr, weight_decay=0.05)
steps = args.epochs * max(1, len(idx["train"]) // args.bs)
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=args.lr, total_steps=steps, pct_start=0.05)


@torch.no_grad()
def retrieval(ids, size=None, reps=50):
    head.eval()
    zi = F.normalize(HE[ids, 0], dim=-1)
    zt = F.normalize(head(TX[ids].to(dev)).float().cpu(), dim=-1)

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

    if size is None or len(ids) <= size:
        return {"he->text": r1(zi, zt), "text->he": r1(zt, zi), "N": len(ids)}
    rng = np.random.default_rng(0)
    a, b = [], []
    for _ in range(reps):
        s = torch.from_numpy(rng.choice(len(ids), size, replace=False))
        a.append(r1(zi[s], zt[s])); b.append(r1(zt[s], zi[s]))
    return {"he->text": float(np.mean(a)), "text->he": float(np.mean(b)), "N": size}


best, hist = -1, []
for ep in range(args.epochs):
    head.train()
    perm = np.random.permutation(idx["train"])
    for s in range(0, len(perm) - args.bs + 1, args.bs):
        b = perm[s:s + args.bs]
        zi = F.normalize(HE[b, torch.randint(0, 8, (len(b),))].to(dev), dim=-1)
        zt = F.normalize(head(TX[b].to(dev)), dim=-1)
        lg = logit.exp().clamp(max=100) * zi @ zt.T
        tgt = torch.arange(len(b), device=dev)
        loss = (F.cross_entropy(lg, tgt) + F.cross_entropy(lg.T, tgt)) / 2
        opt.zero_grad(set_to_none=True); loss.backward(); opt.step(); sched.step()
    v = retrieval(idx["val"])
    vm = (v["he->text"] + v["text->he"]) / 2
    hist.append({"epoch": ep, "loss": loss.item(), "val_r1": vm})
    if vm > best:
        best = vm
        torch.save(head.state_dict(), os.path.join(out_dir, "head.pt"))

head.load_state_dict(torch.load(os.path.join(out_dir, "head.pt")))
res = {"args": vars(args), "best_val_r1": best, "test_full": retrieval(idx["test"]),
       "test_g100": retrieval(idx["test"], size=100), "n": {k: len(v) for k, v in idx.items()}, "history": hist}
# zero-shot baseline on the same test pairs: the text tower's own image tower, no training
zs_img = os.path.join(CACHE, f"pathgen_he_{args.txt_enc}.npy")
if os.path.exists(zs_img):
    zi = F.normalize(torch.from_numpy(np.load(zs_img)[idx["test"], 0].astype(np.float32)), dim=-1)
    zt = F.normalize(TX[idx["test"]], dim=-1)
    s = zi @ zt.T
    res["zeroshot_test_full"] = {"he->text": float(((s > s.diag()[:, None]).sum(1) == 0).float().mean()),
                                 "text->he": float(((s.T > s.diag()[:, None]).sum(1) == 0).float().mean())}
json.dump(res, open(os.path.join(out_dir, "metrics.json"), "w"), indent=1)
print(args.txt_enc, "val", round(best, 3), "test full", res["test_full"], "g100", res["test_g100"])