"""Encoder benchmark for the H&E / mIF / markers / text binding (ImageBind-style, frozen H&E anchor).

With a frozen anchor and anchor-only losses, each trainable encoder receives gradient only from its own
InfoNCE(he, m) term, so one run trains one IF encoder, one marker encoder and one text head *independently*;
each modality keeps its own best epoch (val R@1 of he<->m) and its test embeddings are saved at that epoch.
Cross pairs (if<->markers, if<->text, markers<->text) are then evaluated post hoc for any encoder combination
(10_bench_summary.py).

IF encoders     : resnet18 | convnext_tiny | chanset (shared per-channel ResNet18 + channel tokens + transformer)
marker encoders : linear | mlp | ftt (FT-Transformer, one token per marker)
text encoders   : plip | biomedclip | imagebind   (frozen text tower, trainable MLP head)
"""
import argparse
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 common import CACHE, RUNS, SPLITS, dihedral

p = argparse.ArgumentParser()
p.add_argument("--split", default="orion_official")
p.add_argument("--fold", type=int, default=0)
p.add_argument("--anchor", default="uni2")
p.add_argument("--if_enc", default="resnet18", choices=["resnet18", "convnext_tiny", "chanset"])
p.add_argument("--mk_enc", default="mlp", choices=["linear", "mlp", "ftt"])
p.add_argument("--txt_enc", default="plip", choices=["plip", "biomedclip", "imagebind"])
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("--seed", type=int, default=0)
args = p.parse_args()
torch.manual_seed(args.seed); np.random.seed(args.seed)
dev = "cuda"
TRAIN = ["if", "markers", "text"]
ENC = {"if": args.if_enc, "markers": args.mk_enc, "text": args.txt_enc}
run_dir = os.path.join(RUNS, "bench", f"{args.split}_fold{args.fold}_{args.anchor}_{args.if_enc}_{args.mk_enc}_{args.txt_enc}")
os.makedirs(run_dir, exist_ok=True)

# ---------------- data ----------------
sp = pd.read_csv(os.path.join(SPLITS, args.split, f"fold{args.fold}.csv"))
HE = torch.from_numpy(np.load(os.path.join(CACHE, f"he_{args.anchor}.npy")).astype(np.float32))  # N,8,D
D = HE.shape[-1]
IF = torch.from_numpy(np.load(os.path.join(CACHE, "if.npy")))
MK = torch.from_numpy(np.load(os.path.join(CACHE, "markers.npy")))
TX = torch.from_numpy(np.load(os.path.join(CACHE, f"text_{args.txt_enc}.npy")).astype(np.float32))  # N,4,Dt
tr_idx, va_idx, te_idx = (sp.tile_id[sp.split == s].values for s in ["train", "val", "test"])
if_mean = IF[tr_idx].float().mean((0, 2, 3)).view(1, -1, 1, 1).to(dev) / 255
if_std = IF[tr_idx].float().std((0, 2, 3)).view(1, -1, 1, 1).to(dev) / 255 + 1e-3
mk_mean, mk_std = MK[tr_idx].mean(0), MK[tr_idx].std(0) + 1e-3


def batch(idx, train=False):
    idx = torch.as_tensor(idx)
    d = torch.randint(0, 8, (len(idx),)) if train else torch.zeros(len(idx), dtype=torch.long)
    v = torch.randint(0, 4, (len(idx),)) if train else torch.zeros(len(idx), dtype=torch.long)
    x = IF[idx].to(dev, non_blocking=True).float() / 255
    x = torch.stack([dihedral(a, int(k)) for a, k in zip(x, d)])
    if train:
        x = x * (1 + 0.1 * torch.randn(len(idx), x.shape[1], 1, 1, device=dev))
    mk = (MK[idx] - mk_mean) / mk_std
    if train:
        mk = mk + 0.05 * torch.randn_like(mk)
    return {"he": HE[idx, d].to(dev), "if": (x - if_mean) / if_std, "markers": mk.to(dev), "text": TX[idx, v].to(dev)}


# ---------------- encoders ----------------
def head(din):
    return nn.Sequential(nn.LayerNorm(din), nn.Linear(din, D))


class TimmEnc(nn.Module):
    def __init__(self, name):
        super().__init__()
        tag = {"resnet18": "resnet18.a1_in1k", "convnext_tiny": "convnext_tiny.fb_in1k"}[name]
        self.trunk = timm.create_model(tag, pretrained=True, in_chans=17, num_classes=0)
        self.head = head(self.trunk.num_features)

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


class ChanSet(nn.Module):
    """Channel-agnostic mIF encoder: each channel -> shared 1-ch ResNet18 (112px) -> token (+ channel embedding)
    -> 2-layer transformer over channel tokens -> mean pool."""
    def __init__(self, C=17, d=256):
        super().__init__()
        self.cnn = timm.create_model("resnet18.a1_in1k", pretrained=True, in_chans=1, num_classes=0)
        self.proj = nn.Linear(512, d)
        self.ch = nn.Parameter(torch.zeros(1, C, d)); nn.init.normal_(self.ch, std=0.02)
        self.tf = nn.TransformerEncoder(nn.TransformerEncoderLayer(d, 4, 4 * d, 0.1, batch_first=True, norm_first=True), 2)
        self.head = head(d)

    def forward(self, x):
        B, C, H, W = x.shape
        x = F.interpolate(x, size=(112, 112), mode="bilinear", antialias=True)
        t = self.proj(self.cnn(x.reshape(B * C, 1, 112, 112))).view(B, C, -1) + self.ch
        return self.head(self.tf(t).mean(1))


class FTT(nn.Module):
    """FT-Transformer: feature tokenizer (value * w_j + b_j) + CLS, 3 layers."""
    def __init__(self, F_in, d=128):
        super().__init__()
        self.w = nn.Parameter(torch.randn(F_in, d) * 0.02)
        self.b = nn.Parameter(torch.zeros(F_in, d))
        self.cls = nn.Parameter(torch.zeros(1, 1, d))
        self.tf = nn.TransformerEncoder(nn.TransformerEncoderLayer(d, 8, 2 * d, 0.1, batch_first=True, norm_first=True), 3)
        self.head = head(d)

    def forward(self, x):
        t = x.unsqueeze(-1) * self.w + self.b
        t = torch.cat([self.cls.expand(len(x), -1, -1), t], 1)
        return self.head(self.tf(t)[:, 0])


def mlp(din, hid=512):
    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, D))


enc = nn.ModuleDict({
    "if": ChanSet() if args.if_enc == "chanset" else TimmEnc(args.if_enc),
    "markers": {"linear": lambda: nn.Linear(MK.shape[1], D), "mlp": lambda: mlp(MK.shape[1]),
                "ftt": lambda: FTT(MK.shape[1])}[args.mk_enc](),
    "text": mlp(TX.shape[-1]),  # frozen text tower features -> trainable projection
}).to(dev)
logit = nn.ParameterDict({m: nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32)) for m in TRAIN}).to(dev)
params = list(enc.parameters()) + list(logit.parameters())
opt = torch.optim.AdamW(params, lr=args.lr, weight_decay=0.05)
steps = args.epochs * (len(tr_idx) // args.bs)
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=args.lr, total_steps=steps, pct_start=0.05)


def embed(b):
    z = {"he": F.normalize(b["he"], dim=-1)}
    for m in TRAIN:
        z[m] = F.normalize(enc[m](b[m]).float(), dim=-1)
    return z


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


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


# ---------------- train with per-modality model selection ----------------
best = {m: -1.0 for m in TRAIN}
best_ep = {m: -1 for m in TRAIN}
test_z = {}
hist = []
t0 = time.time()
for ep in range(args.epochs):
    enc.train()
    perm = np.random.permutation(tr_idx)
    tl = {m: [] for m in TRAIN}
    for s in range(0, len(perm) - args.bs + 1, args.bs):
        b = batch(perm[s:s + args.bs], train=True)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            z = embed(b)
        tgt = torch.arange(args.bs, device=dev)
        loss = 0
        for m in TRAIN:
            lg = logit[m].exp().clamp(max=100) * z["he"] @ z[m].T
            l = (F.cross_entropy(lg, tgt) + F.cross_entropy(lg.T, tgt)) / 2
            loss = loss + l
            tl[m].append(l.item())
        opt.zero_grad(set_to_none=True)
        loss.backward()
        opt.step(); sched.step()
    zv = embed_all(va_idx)
    rec = {"epoch": ep}
    improved = []
    for m in TRAIN:
        v = (r1(zv["he"], zv[m]) + r1(zv[m], zv["he"])) / 2
        rec[f"loss_{m}"], rec[f"val_r1_{m}"] = float(np.mean(tl[m])), v
        if v > best[m]:
            best[m], best_ep[m] = v, ep
            improved.append(m)
    if improved:
        zt = embed_all(te_idx)
        test_z["he"] = zt["he"]
        for m in improved:
            test_z[m] = zt[m]
    hist.append(rec)
    print(f"ep {ep} " + " ".join(f"{m}:{rec[f'val_r1_{m}']:.3f}" for m in TRAIN) + f" ({time.time() - t0:.0f}s)", flush=True)

torch.save({"z": test_z, "tile_id": te_idx}, os.path.join(run_dir, "test_embeddings.pt"))
json.dump({"args": vars(args), "best_val": best, "best_epoch": best_ep, "n_train": len(tr_idx), "n_test": len(te_idx),
           "time_s": time.time() - t0, "history": hist}, open(os.path.join(run_dir, "train.json"), "w"), indent=1)
print("done", best, best_ep)