"""MIPHEI-ViT (Sanofi, arXiv 2505.10294) virtual mIF: H&E -> 16 ORION channels (all ORION markers except PD-1).

Model : ext/miphei/weights (GitHub release v1.0.0: decoder + LoRA, 26 MB; non-commercial academic licence) on top of
        bioptimus/H-optimus-0 (HF, gated; HF token only via environment). model.py of the release (read in full) is
        loaded from its path; nothing else from the downloaded repo is executed.
Training data of MIPHEI = this very ORION tile set with the ORION official split (37 train / 2 val / 2 test slides =
        our splits/orion_official) -> MIPHEI virtual mIF is held-out ONLY on the official test (and val) tiles.
Input : RGB at 0.5 um/px (ORION tiles 333 px = 0.4997 um/px used as is; TCGA 672-px patches at ~0.25 um/px -> 336 px),
        ImageNet mean/std (release config), reflect-padded to 512 px (input must be a power of 2), output cropped back.
Output: tanh in [-0.9, 0.9] -> (y + 0.9) / 1.8 * 255 = the 8-bit scale of the real ORION mIF tiles (MIPHEI targets).
  cache/miphei/<set>_maps168.npy uint8 (N, 16, 168, 168)  resampled exactly like the real tiles in 14_ (antialiased
                                                           bilinear to ~1 um/px) -> VirTues input
  cache/miphei/<set>_means.npy   float32 (N, 16)          tile means at full resolution (0-255 scale)
  cache/miphei/<set>_meta.csv, channels.json
Usage: HF_TOKEN=... .venv/bin/python 35_miphei_infer.py --set orion|tcga_text|dlpfc|idc [--bs 8]
"""
import argparse
import importlib.util
import json
import os
import time

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image, PngImagePlugin
from safetensors.torch import load_file

from common import CACHE, ORION, ROOT, load_all_tiles

p = argparse.ArgumentParser()
p.add_argument("--set", required=True)
p.add_argument("--bs", type=int, default=8)
args = p.parse_args()
dev = "cuda"
W = os.path.join(ROOT, "ext", "miphei", "weights")
OUT = os.path.join(CACHE, "miphei")
os.makedirs(OUT, exist_ok=True)
CH = ["Hoechst", "CD31", "CD45", "CD68", "CD4", "FOXP3", "CD8a", "CD45RO", "CD20", "PD-L1", "CD3e", "CD163",
      "E-cadherin", "Ki67", "Pan-CK", "SMA"]  # release config targ_channel_names (order of the output heads)
MEAN = torch.tensor([123.675, 116.28, 103.53], device=dev).view(1, 3, 1, 1) / 255
STD = torch.tensor([58.395, 57.12, 57.375], device=dev).view(1, 3, 1, 1) / 255
PAD_TO, VSIZE = 512, 168
PngImagePlugin.MAX_TEXT_CHUNK = 64 * 1024 ** 2  # some TCGA PNGs carry large text metadata chunks

spec = importlib.util.spec_from_file_location("miphei_model", os.path.join(W, "model.py"))
Mm = importlib.util.module_from_spec(spec); spec.loader.exec_module(Mm)
vit = Mm.get_hoptimus0_hf("bioptimus/H-optimus-0")
vit.set_input_size(img_size=(PAD_TO, PAD_TO))
enc = Mm.Encoder(vit)
model = Mm.MIPHEIViT(enc, Mm.Detail_Capture(emb_chans=enc.embed_dim, out_chans=len(CH), use_attention=True, activation=nn.Tanh()))
sd = Mm.merge_lora_weights(model, load_file(os.path.join(W, "model.safetensors")))
Mm.validate_load_info(model.load_state_dict(sd, strict=False))
model = model.eval().to(dev)

if args.set == "orion":
    df = load_all_tiles()
    paths, keys, scale = [os.path.join(ORION, x) for x in df.image_path], df.image_path.values, 1.0
elif args.set == "tcga_text":
    I = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv"))
    paths, keys, scale = list(I.image_path), I.key.values, 0.5
elif args.set in ("dlpfc", "idc"):  # external sets (36_external_prep.py): uint8 RGB patches at 0.5 um/px, 224 px
    E = os.path.join(CACHE, "external")
    ARR = np.load(os.path.join(E, "dlpfc_he_05um.npy" if args.set == "dlpfc" else "idc_he.npy"), mmap_mode="r")
    em = pd.read_csv(os.path.join(E, f"{args.set}_meta.csv"))
    assert len(em) == len(ARR)
    paths, keys, scale = list(range(len(ARR))), em.index.values, 1.0
else:
    raise ValueError(args.set)
N = len(paths)
maps = np.lib.format.open_memmap(os.path.join(OUT, f"{args.set}_maps168.npy"), "w+", np.uint8, (N, len(CH), VSIZE, VSIZE))
means = np.zeros((N, len(CH)), np.float32)
sizes = []
t0 = time.time()


def load(pth):
    if isinstance(pth, int):
        return np.asarray(ARR[pth])
    im = Image.open(pth).convert("RGB")
    if scale != 1.0:
        im = im.resize((round(im.width * scale), round(im.height * scale)), Image.BICUBIC)
    return np.asarray(im)


with torch.inference_mode():
    for s in range(0, N, args.bs):
        ims = [load(x) for x in paths[s:s + args.bs]]
        h, w = ims[0].shape[:2]
        assert all(i.shape[:2] == (h, w) for i in ims) and h <= PAD_TO and w <= PAD_TO
        x = torch.from_numpy(np.stack(ims)).to(dev).permute(0, 3, 1, 2).float() / 255
        x = (x - MEAN) / STD
        pt, pl = (PAD_TO - h) // 2, (PAD_TO - w) // 2
        x = F.pad(x, (pl, PAD_TO - w - pl, pt, PAD_TO - h - pt), mode="reflect")
        with torch.autocast("cuda", dtype=torch.float16):  # as in the release notebook
            y = model(x)
        y = y.float().clamp(-0.9, 0.9)[:, :, pt:pt + h, pl:pl + w]
        y = (y + 0.9) / 1.8 * 255
        means[s:s + len(ims)] = y.mean((2, 3)).cpu().numpy()
        maps[s:s + len(ims)] = F.interpolate(y, size=(VSIZE, VSIZE), mode="bilinear", antialias=True,
                                             align_corners=False).round().clamp(0, 255).byte().cpu().numpy()
        sizes += [(h, w)] * len(ims)
        if s % (args.bs * 100) == 0:
            print(f"  {s}/{N} [{time.time() - t0:.0f}s]", flush=True)
maps.flush()
np.save(os.path.join(OUT, f"{args.set}_means.npy"), means)
pd.DataFrame({"key": keys, "path": paths, "h": [a for a, _ in sizes], "w": [b for _, b in sizes]}).to_csv(
    os.path.join(OUT, f"{args.set}_meta.csv"), index=False)
json.dump({"channels": CH}, open(os.path.join(OUT, "channels.json"), "w"))
print(f"done {args.set}: {N} tiles [{time.time() - t0:.0f}s]")