"""MedVTok (general medical visual tokenizer, MERA-aligned to CONCH v1.5 for histopathology) global image features
of every H&E set, as an additional / alternative image encoder next to the frozen UNI2 anchor.

Run:  .venv/bin/python -I 50_medvtok_feats.py <set> [<set> ...]      sets: orion tcga hest dlpfc
Model: ext/medvtok/repo (Masaaki-75/medvtok @a558c24, MIT; only models/medvtok.py + models/utils.py +
models/transformer.py are imported) + ext/medvtok/weights/ckpt.pth (HF massaki75/medvtok, MIT), loaded here with
torch.load(weights_only=True) and strict=True (the repo's build_model uses weights_only=False / strict=False).
Encoder only (decode=False). Preprocessing = repo build_transform: RGB -> bicubic resize 256 -> [-1, 1].
Feature = outputs["feats"]: mean-pooled encoder tokens -> LayerNorm -> Linear (1152-d, the aligned global feature).
Views: orion / tcga 8 dihedral transforms (= he_uni2.npy); hest / dlpfc 1 view (= their UNI2 caches).
Images: orion ../ORION tiles, tcga cache/allimg_meta.csv image_path, hest cache/hest/he_img.npy (49_),
dlpfc cache/external/dlpfc_he_05um.npy (0.5 um/px, as for MIPHEI / captions).
Output: cache/he_medvtok.npy (3187, 8, 1152), cache/allimg_he_medvtok.npy (6144, 8, 1152),
cache/hest/he_medvtok.npy (9484, 1, 1152), cache/external/dlpfc_he_medvtok.npy (482, 1152); float16.
"""
import json
import os
import sys

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

PngImagePlugin.MAX_TEXT_CHUNK = 64 * 1024 * 1024  # some PathGen PNGs embed very large ICC profiles (= 04b_)

ROOT = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, ROOT)
from common import CACHE, ORION, dihedral, load_all_tiles  # noqa: E402

REPO = os.path.join(ROOT, "ext", "medvtok", "repo")
WTS = os.path.join(ROOT, "ext", "medvtok", "weights")
sys.path.insert(0, REPO)
from models.medvtok import MedVTok  # noqa: E402

dev = "cuda"
SIZE, BS = 256, 64

cfg = json.load(open(os.path.join(WTS, "config.json")))
assert cfg["architecture"] == "medvtok-clip-meanpool"
model = MedVTok(**cfg["model_kwargs"])
sd = torch.load(os.path.join(WTS, "ckpt.pth"), map_location="cpu", weights_only=True)
print(model.load_state_dict(sd, strict=True))
del sd, model.decoder
model.eval().to(dev)


@torch.no_grad()
def encode(u8):
    """u8: (B, H, W, 3) uint8 -> (B, 1152) float32"""
    x = torch.from_numpy(np.ascontiguousarray(u8)).to(dev).permute(0, 3, 1, 2).float() / 255
    x = F.interpolate(x, size=(SIZE, SIZE), mode="bicubic", align_corners=False, antialias=True).clamp(0, 1)
    x = x * 2 - 1
    with torch.autocast("cuda", dtype=torch.bfloat16):
        out = model([im for im in x], sample_posterior=False, decode=False, return_batched_tensor=False)
    return out["feats"].float().cpu().numpy()


def run(load, n, views):
    out = np.zeros((n, views, 1152), np.float16)
    for s in range(0, n, BS):
        u8 = load(s, min(n, s + BS))
        for d in range(views):
            x = u8 if views == 1 else dihedral(torch.from_numpy(u8).permute(0, 3, 1, 2), d).permute(0, 2, 3, 1).numpy()
            out[s:s + len(u8), d] = encode(x)
        if s % (BS * 20) == 0:
            print(f"  {s}/{n}", flush=True)
    assert np.isfinite(out.astype(np.float32)).all()
    return out


def from_paths(paths):
    return lambda a, b: np.stack([np.asarray(Image.open(p).convert("RGB").resize((SIZE, SIZE), Image.BICUBIC))
                                  for p in paths[a:b]])


for st in sys.argv[1:]:
    if st == "orion":
        paths = np.array([os.path.join(ORION, p) for p in load_all_tiles().image_path])
        np.save(os.path.join(CACHE, "he_medvtok.npy"), run(from_paths(paths), len(paths), 8))
    elif st == "tcga":
        paths = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv")).image_path.values
        np.save(os.path.join(CACHE, "allimg_he_medvtok.npy"), run(from_paths(paths), len(paths), 8))
    elif st == "hest":
        A = np.load(os.path.join(CACHE, "hest", "he_img.npy"), mmap_mode="r")
        np.save(os.path.join(CACHE, "hest", "he_medvtok.npy"), run(lambda a, b: np.asarray(A[a:b]), len(A), 1))
    elif st == "dlpfc":
        A = np.load(os.path.join(CACHE, "external", "dlpfc_he_05um.npy"))
        np.save(os.path.join(CACHE, "external", "dlpfc_he_medvtok.npy"), run(lambda a, b: A[a:b], len(A), 1)[:, 0])
    print("saved", st, flush=True)