"""Bind the gene-expression modality (HEST Visium spots) to the frozen H&E anchor (ImageBind recipe).

anchor (frozen, no head): UNI2-h (same model as the ORION / PathGen anchor) or GigaPath spot features
gene encoders:  mlp            MLP on log1p-CP10k of the 2000 HVGs (from scratch)
                fm:scgpt       frozen scGPT-human spot embedding (22_) + MLP head
                fm:geneformer  frozen Geneformer-V2 spot embedding (22_) + MLP head
loss: symmetric InfoNCE(he, gene); the gene head maps into the raw anchor space, so the same space also holds the
ORION mIF head and the PathGen text head (emergent gene<->protein, gene<->text evaluated in 24_).
CV: 5-fold GroupKFold over study+patient groups (21_), 10% of training groups held out for early stopping.
Outputs runs/gene/<anchor>_<enc>_fold<k>/{head.pt, test.pt}, results/gene_retrieval.csv
"""
import argparse
import os

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.model_selection import GroupKFold

from common import CACHE, ROOT, RUNS, metrics

p = argparse.ArgumentParser()
p.add_argument("--anchors", nargs="+", default=["uni2", "gigapath"])
p.add_argument("--encs", nargs="+", default=["mlp", "fm:scgpt", "fm:geneformer"])
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)
args = p.parse_args()
dev = "cuda"
H = os.path.join(CACHE, "hest")
M = pd.read_csv(os.path.join(H, "meta.csv"))


def feats(enc):
    f = os.path.join(H, "hvg.npy" if enc == "mlp" else f"gene_{enc.split(':')[1]}.npy")
    x = np.load(f).astype(np.float32)
    return torch.from_numpy(x.reshape(len(x), -1)).to(dev)


def mlp(din, dout, hid=512):  # same head as 09_bench / 12_ / 20_
    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 r1(a, b):
    s = a @ b.T
    return float(((s > s.diag()[:, None]).sum(1) == 0).float().mean())


def train(A, G, tr, va, seed=0):
    torch.manual_seed(seed); rng = np.random.default_rng(seed)
    head = mlp(G.shape[1], A.shape[1]).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(tr) // args.bs)
    sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=args.lr, total_steps=steps, pct_start=0.05)
    za_va = F.normalize(A[va], dim=-1)
    best, state = -1, None
    for ep in range(args.epochs):
        head.train()
        perm = rng.permutation(tr)
        for s in range(0, len(perm) - args.bs + 1, args.bs):
            b = torch.as_tensor(perm[s:s + args.bs], device=dev)
            za, zg = F.normalize(A[b], dim=-1), F.normalize(head(G[b]), dim=-1)
            lg = logit.exp().clamp(max=100) * za @ zg.T
            t = torch.arange(len(b), device=dev)
            loss = (F.cross_entropy(lg, t) + F.cross_entropy(lg.T, t)) / 2
            opt.zero_grad(set_to_none=True); loss.backward(); opt.step(); sched.step()
        head.eval()
        with torch.no_grad():
            zg = F.normalize(head(G[va]), dim=-1)
            v = (r1(za_va, zg) + r1(zg, za_va)) / 2
        if v > best:
            best, state = v, {k: x.clone() for k, x in head.state_dict().items()}
    head.load_state_dict(state); head.eval()
    return head, best


folds = list(GroupKFold(5).split(M, groups=M.group))
rows = []
for an in args.anchors:
    A = torch.from_numpy(np.load(os.path.join(H, f"he_{an}.npy"))[:, 0].astype(np.float32)).to(dev)
    for enc in args.encs:
        G = feats(enc)
        for k, (trv, te) in enumerate(folds):
            grp = np.random.default_rng(k).permutation(sorted(M.group.iloc[trv].unique()))
            va_g = set(grp[: max(1, len(grp) // 10)])
            is_va = M.group.iloc[trv].isin(va_g).values
            tr, va = trv[~is_va], trv[is_va]
            head, best = train(A, G, tr, va)
            with torch.no_grad():
                za, zg = F.normalize(A[te], dim=-1).cpu(), F.normalize(head(G[te]), dim=-1).cpu()
            d = os.path.join(RUNS, "gene", f"{an}_{enc.replace(':', '-')}_fold{k}")
            os.makedirs(d, exist_ok=True)
            torch.save(head.state_dict(), os.path.join(d, "head.pt"))
            torch.save({"idx": te, "he": za, "gene": zg}, os.path.join(d, "test.pt"))
            for q, g, name in [(za, zg, "he->gene"), (zg, za, "gene->he")]:
                rows.append({"anchor": an, "gene_enc": enc, "fold": k, "dir": name, "val_r1": best, **metrics(q, g)})
            print(f"{an} {enc} fold{k}: val {best:.3f}  test R@1_g100 {np.mean([r['R@1_g100'] for r in rows[-2:]]):.3f} "
                  f"pair_acc {np.mean([r['pair_acc'] for r in rows[-2:]]):.3f} (n_test={len(te)})", flush=True)
        pd.DataFrame(rows).to_csv(os.path.join(ROOT, "results", "gene_retrieval.csv"), index=False)
R = pd.DataFrame(rows)
g = R.groupby(["anchor", "gene_enc", "fold"])[["R@1", "R@1_g100", "R@10", "pair_acc"]].mean()
print(g.groupby(["anchor", "gene_enc"]).agg(["mean", "std"]).round(3).to_string())