"""ImageBind-style binding of IF / nuclei / markers to a frozen H&E (ImageBind-Huge vision) anchor.

loss = sum_m InfoNCE_sym(he, m)       (--loss anchor, ImageBind default; IF<->nuclei etc. are *emergent*)
     + sum_{m<n} InfoNCE_sym(m, n)    (--loss allpairs ablation)

Evaluation on the test split: cross-modal retrieval for all 12 directed modality pairs
(R@1/5/10, MRR, median rank) and one-to-one pairing accuracy (Hungarian assignment on similarity).
"""
import argparse
import itertools
import json
import os
import time

import numpy as np
import pandas as pd
import timm
import torch
import torch.nn as nn
import torch.nn.functional as F
from scipy.optimize import linear_sum_assignment

from common import CACHE, MODALITIES, RUNS, SPLITS, dihedral

p = argparse.ArgumentParser()
p.add_argument("--split", default="orion_official")
p.add_argument("--fold", type=int, default=0)
p.add_argument("--loss", default="anchor", choices=["anchor", "allpairs"])
p.add_argument("--anchor", default="imagebind", choices=["imagebind", "uni2", "conch", "gigapath", "plip", "phikonv2"],
               help="frozen H&E encoder whose features (cache/he_<anchor>.npy) define the shared space")
p.add_argument("--epochs", type=int, default=40)
p.add_argument("--bs", type=int, default=128)
p.add_argument("--lr", type=float, default=3e-4)
p.add_argument("--backbone", default="resnet18.a1_in1k")
p.add_argument("--seed", type=int, default=0)
p.add_argument("--min_nuclei_tissue", type=int, default=20)
args = p.parse_args()
torch.manual_seed(args.seed); np.random.seed(args.seed)
dev = "cuda"
run_dir = os.path.join(RUNS, f"{args.split}_fold{args.fold}_{args.loss}"
                       + ("" if args.anchor == "imagebind" else f"_{args.anchor}"))
os.makedirs(run_dir, exist_ok=True)

# ---------------- data ----------------
sp = pd.read_csv(os.path.join(SPLITS, args.split, f"fold{args.fold}.csv"))
he_emb = torch.from_numpy(np.load(os.path.join(CACHE, f"he_{args.anchor}.npy")).astype(np.float32))  # N,8,D
D = he_emb.shape[-1]  # shared embedding dim = anchor dim (1024 ImageBind/Phikon-v2, 1536 UNI2/GigaPath, 512 CONCH/PLIP)
IF = torch.from_numpy(np.load(os.path.join(CACHE, "if.npy")))          # N,17,224,224 uint8 (in RAM, ~2.7GB)
NUC = torch.from_numpy(np.load(os.path.join(CACHE, "nuclei.npy")))     # N,2,224,224
MK = torch.from_numpy(np.load(os.path.join(CACHE, "markers.npy")))     # N,33
tr_idx = sp.tile_id[sp.split == "train"].values
va_idx = sp.tile_id[sp.split == "val"].values
te_idx = sp.tile_id[sp.split == "test"].values
# normalisation statistics from training tiles only
if_mean = IF[tr_idx].float().mean((0, 2, 3)).view(1, -1, 1, 1) / 255
if_std = IF[tr_idx].float().std((0, 2, 3)).view(1, -1, 1, 1) / 255 + 1e-3
mk_mean, mk_std = MK[tr_idx].mean(0), MK[tr_idx].std(0) + 1e-3


def batch(idx, d=None, train=False):
    """Returns dict of modality tensors on GPU. d: dihedral id per sample (shared across image modalities)."""
    idx = torch.as_tensor(idx)
    if d is None:
        d = torch.zeros(len(idx), dtype=torch.long)
    x_if = IF[idx].to(dev, non_blocking=True).float() / 255
    x_nu = NUC[idx].to(dev, non_blocking=True).float() / 255
    x_if = torch.stack([dihedral(a, int(k)) for a, k in zip(x_if, d)])
    x_nu = torch.stack([dihedral(a, int(k)) for a, k in zip(x_nu, d)])
    if train:  # mild per-channel intensity jitter for IF
        x_if = x_if * (1 + 0.1 * torch.randn(len(idx), x_if.shape[1], 1, 1, device=dev))
    x_if = (x_if - if_mean.to(dev)) / if_std.to(dev)
    x_nu = (x_nu - 0.5) / 0.5
    mk = ((MK[idx] - mk_mean) / mk_std).to(dev)
    if train:
        mk = mk + 0.05 * torch.randn_like(mk)
    he = he_emb[idx, d].to(dev)
    return {"he": he, "if": x_if, "nuclei": x_nu, "markers": mk}


# ---------------- model ----------------
class ImgEnc(nn.Module):
    def __init__(self, in_chans):
        super().__init__()
        self.trunk = timm.create_model(args.backbone, pretrained=True, in_chans=in_chans, num_classes=0)
        self.head = nn.Sequential(nn.LayerNorm(self.trunk.num_features), nn.Linear(self.trunk.num_features, D))

    def forward(self, x):
        return self.head(self.trunk(x))


class MLPEnc(nn.Module):
    def __init__(self, d_in):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(d_in, 512), nn.GELU(), nn.Dropout(0.1), nn.Linear(512, 512), nn.GELU(),
                                 nn.LayerNorm(512), nn.Linear(512, D))

    def forward(self, x):
        return self.net(x)


class Bind(nn.Module):
    def __init__(self):
        super().__init__()
        self.enc = nn.ModuleDict({"if": ImgEnc(17), "nuclei": ImgEnc(2), "markers": MLPEnc(MK.shape[1])})
        pairs = [("he", m) for m in ["if", "nuclei", "markers"]]
        if args.loss == "allpairs":
            pairs += [("if", "nuclei"), ("if", "markers"), ("nuclei", "markers")]
        self.pairs = pairs
        # ImageBind: learnable logit scale (init 1/0.07) per pair
        self.logit_scale = nn.ParameterDict({f"{a}-{b}": nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32))
                                             for a, b in pairs})

    def embed(self, b):
        z = {"he": F.normalize(b["he"], dim=-1)}  # frozen anchor
        for m, e in self.enc.items():
            z[m] = F.normalize(e(b[m]), dim=-1)
        return z

    def loss(self, z):
        tot, logs = 0, {}
        tgt = torch.arange(len(z["he"]), device=dev)
        for a, bm in self.pairs:
            s = self.logit_scale[f"{a}-{bm}"].exp().clamp(max=100)
            lg = s * z[a] @ z[bm].T
            l = (F.cross_entropy(lg, tgt) + F.cross_entropy(lg.T, tgt)) / 2
            tot = tot + l
            logs[f"{a}-{bm}"] = l.item()
        return tot, logs


model = Bind().to(dev)
opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.05)
steps = args.epochs * max(1, len(tr_idx) // args.bs)
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=args.lr, total_steps=steps, pct_start=0.05)
scaler = torch.amp.GradScaler()


@torch.no_grad()
def embed_all(idx):
    model.eval()
    zs = {m: [] for m in MODALITIES}
    for s in range(0, len(idx), 256):
        with torch.autocast("cuda", dtype=torch.bfloat16):
            z = model.embed(batch(idx[s:s + 256]))
        for m in MODALITIES:
            zs[m].append(z[m].float().cpu())
    return {m: torch.cat(v) for m, v in zs.items()}


def retrieval_metrics(zq, zg):
    sim = (zq @ zg.T).numpy()
    n = len(sim)
    rank = (sim > sim[np.arange(n), np.arange(n)][:, None]).sum(1)  # 0 = correct is top-1
    r, c = linear_sum_assignment(-sim)
    return {"R@1": float((rank < 1).mean()), "R@5": float((rank < 5).mean()), "R@10": float((rank < 10).mean()),
            "MRR": float((1 / (rank + 1)).mean()), "MedR": float(np.median(rank + 1)),
            "pair_acc": float((c[np.argsort(r)] == np.arange(n)).mean()), "N": n, "chance_R@1": 1 / n}


def fixed_gallery(zq, zg, size=100, reps=50):
    """R@1 / pairing accuracy averaged over random galleries of a fixed size (comparable across splits)."""
    rng = np.random.default_rng(0)
    n = len(zq)
    if n <= size:
        m = retrieval_metrics(zq, zg)
        return {"R@1_g100": m["R@1"], "pair_acc_g100": m["pair_acc"]}
    r1, pa = [], []
    for _ in range(reps):
        s = torch.from_numpy(rng.choice(n, size, replace=False))
        m = retrieval_metrics(zq[s], zg[s])
        r1.append(m["R@1"]); pa.append(m["pair_acc"])
    return {"R@1_g100": float(np.mean(r1)), "pair_acc_g100": float(np.mean(pa))}


def evaluate(idx, fixed=False):
    z = embed_all(idx)
    res = {}
    for a, b in itertools.permutations(MODALITIES, 2):
        res[f"{a}->{b}"] = retrieval_metrics(z[a], z[b])
        if fixed:
            res[f"{a}->{b}"].update(fixed_gallery(z[a], z[b]))
    return res, z


def mean_r1(res):
    return float(np.mean([v["R@1"] for v in res.values()]))


# ---------------- train ----------------
best, hist = -1, []
t0 = time.time()
for ep in range(args.epochs):
    model.train()
    perm = np.random.permutation(tr_idx)
    tl = []
    for s in range(0, len(perm) - args.bs + 1, args.bs):
        idx = perm[s:s + args.bs]
        b = batch(idx, d=torch.randint(0, 8, (len(idx),)), train=True)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            loss, logs = model.loss(model.embed(b))
        opt.zero_grad(set_to_none=True)
        loss.backward()
        opt.step(); sched.step()
        tl.append(loss.item())
    vres, _ = evaluate(va_idx)
    v = mean_r1(vres)
    hist.append({"epoch": ep, "train_loss": float(np.mean(tl)), "val_meanR1": v, **{k: x for k, x in logs.items()}})
    print(f"ep {ep} loss {np.mean(tl):.3f} val meanR@1 {v:.3f} ({time.time() - t0:.0f}s)", flush=True)
    if v > best:
        best = v
        torch.save(model.state_dict(), os.path.join(run_dir, "best.pt"))

# ---------------- test ----------------
model.load_state_dict(torch.load(os.path.join(run_dir, "best.pt")))
tres, z = evaluate(te_idx, fixed=True)
stats = pd.read_csv(os.path.join(os.path.dirname(RUNS), "results", "tile_stats.csv")).set_index("tile_id")
tissue = stats.loc[te_idx, "n_nuclei"].values >= args.min_nuclei_tissue
tres_tissue = {}
for a, b in itertools.permutations(MODALITIES, 2):
    tres_tissue[f"{a}->{b}"] = retrieval_metrics(z[a][tissue], z[b][tissue])
    tres_tissue[f"{a}->{b}"].update(fixed_gallery(z[a][tissue], z[b][tissue]))
out = {"args": vars(args), "n_train": len(tr_idx), "n_val": len(va_idx), "n_test": len(te_idx),
       "n_test_tissue": int(tissue.sum()), "best_val_meanR1": best, "test": tres, "test_tissue": tres_tissue,
       "history": hist}
json.dump(out, open(os.path.join(run_dir, "metrics.json"), "w"), indent=1)
torch.save({m: v for m, v in z.items()}, os.path.join(run_dir, "test_embeddings.pt"))
print("TEST mean R@1 (all pairs):", mean_r1(tres), "| tissue-only:", mean_r1(tres_tissue))
for k in ["he->if", "he->nuclei", "he->markers", "if->nuclei", "if->markers", "nuclei->markers"]:
    print(f"  {k:18s} R@1 {tres[k]['R@1']:.3f} R@5 {tres[k]['R@5']:.3f} pair_acc {tres[k]['pair_acc']:.3f} R@1_g100 {tres[k]['R@1_g100']:.3f}")