"""Frozen pathology foundation-model embeddings of every H&E tile under the 8 dihedral transforms.
Usage: python 04b_extract_anchor_fm.py {uni2|conch|conchclip|gigapath|plip|phikonv2} [--set <name>]
Output: cache/he_<name>.npy (N, 8, D) float16 (raw features; L2-normalised later in training).

  uni2     : MahmoodLab/UNI2-h, ViT-H/14 (reg tokens), 224px, ImageNet norm, CLS -> 1536-d
  gigapath : prov-gigapath/prov-gigapath tile encoder, ViT-g/14, 224px, ImageNet norm -> 1536-d
  conch    : MahmoodLab/CONCH v1 vision tower (ViT-B/16), 448px, OpenAI-CLIP norm,
             encode_image(proj_contrast=False, normalize=False) -> 512-d (authors' recommended tile feature)
  plip     : vinid/plip CLIP ViT-B/32 image features (projected) -> 512-d      [open stand-in for CONCH]
  phikonv2 : owkin/phikon-v2 ViT-L/16 DINOv2, CLS -> 1024-d                  [open stand-in for GigaPath]
"""
import os
import sys

import numpy as np
import timm
import torch
from PIL import Image, PngImagePlugin

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

from common import CACHE, ORION, dihedral, load_all_tiles, load_imagebind_vision

name = sys.argv[1]
SET = sys.argv[sys.argv.index("--set") + 1] if "--set" in sys.argv else "pathgen" if "--pathgen" in sys.argv else None
PATHGEN = SET is not None  # encode an image-text pair set (cache/<SET>_meta.csv, e.g. PathGen-1.6M) instead of ORION tiles
UNI2_LOCAL = "/diskarray/home/tl688/reg2026/model_alldata/uni2-h.bin"
dev = "cuda"
IMNET = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
CLIPN = ([0.48145466, 0.4578275, 0.40821073], [0.26862954, 0.26130258, 0.27577711])

if name == "uni2":
    # HF repo is gated for this account -> use the local copy of the official UNI2-h checkpoint
    model = timm.create_model(
        "vit_giant_patch14_224", pretrained=False, img_size=224, patch_size=14, depth=24, num_heads=24,
        init_values=1e-5, embed_dim=1536, mlp_ratio=2.66667 * 2, num_classes=0, no_embed_class=True,
        mlp_layer=timm.layers.SwiGLUPacked, act_layer=torch.nn.SiLU, reg_tokens=8, dynamic_img_size=True)
    model.load_state_dict(torch.load(UNI2_LOCAL, map_location="cpu"), strict=True)
    size, norm, fwd = 224, IMNET, lambda m, x: m(x)
elif name == "gigapath":
    model = timm.create_model("hf_hub:prov-gigapath/prov-gigapath", pretrained=True)
    size, norm, fwd = 224, IMNET, lambda m, x: m(x)
elif name == "conch":
    from conch.open_clip_custom import create_model_from_pretrained
    model, _ = create_model_from_pretrained("conch_ViT-B-16", "hf_hub:MahmoodLab/conch")
    size, norm = 448, CLIPN
    fwd = lambda m, x: m.encode_image(x, proj_contrast=False, normalize=False)
elif name == "conchclip":  # CONCH image tower projected into its joint image-text space (zero-shot with CONCH text)
    from conch.open_clip_custom import create_model_from_pretrained
    model, _ = create_model_from_pretrained("conch_ViT-B-16", "hf_hub:MahmoodLab/conch")
    size, norm = 448, CLIPN
    fwd = lambda m, x: m.encode_image(x, proj_contrast=True, normalize=False)
elif name == "plip":  # open pathology CLIP (OpenPath tweets), stand-in for gated CONCH
    from transformers import CLIPModel
    model = CLIPModel.from_pretrained("vinid/plip")
    size, norm, fwd = 224, CLIPN, lambda m, x: m.get_image_features(pixel_values=x)
elif name == "phikonv2":  # open DINOv2 ViT-L pathology SSL model, stand-in for gated GigaPath
    from transformers import AutoModel
    model = AutoModel.from_pretrained("owkin/phikon-v2")
    size, norm, fwd = 224, IMNET, lambda m, x: m(pixel_values=x).last_hidden_state[:, 0]
elif name == "biomedclip":  # BiomedCLIP ViT-B/16 image tower (PMC-15M), projected -> 512; zero-shot baselines
    import open_clip
    model, _ = open_clip.create_model_from_pretrained("hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224")
    size, norm, fwd = 224, CLIPN, lambda m, x: m.encode_image(x)
elif name == "imagebind":  # used here for PathGen zero-shot baselines
    model, MT = load_imagebind_vision()
    size, norm, fwd = 224, CLIPN, lambda m, x: m({MT.VISION: x})[MT.VISION]
else:
    raise ValueError(name)
model.eval().to(dev)
mean = torch.tensor(norm[0]).view(1, 3, 1, 1)
std = torch.tensor(norm[1]).view(1, 3, 1, 1)

if PATHGEN:
    import pandas as pd
    paths = pd.read_csv(os.path.join(CACHE, f"{SET}_meta.csv")).image_path.values
else:
    paths = np.array([os.path.join(ORION, p) for p in load_all_tiles().image_path])
out = None
bs = 32
with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16):
    for s in range(0, len(paths), bs):
        ims = [np.asarray(Image.open(p).convert("RGB").resize((size, size), Image.BICUBIC))
               for p in paths[s:s + bs]]
        x = (torch.from_numpy(np.stack(ims)).permute(0, 3, 1, 2).float() / 255 - mean) / std
        for d in range(8):
            e = fwd(model, dihedral(x, d).to(dev)).float().cpu().numpy()
            if out is None:
                out = np.zeros((len(paths), 8, e.shape[1]), np.float16)
            out[s:s + bs, d] = e
        if s % 640 == 0:
            print(name, s, len(paths), flush=True)
np.save(os.path.join(CACHE, (f"{SET}_" if PATHGEN else "") + f"he_{name}.npy"), out)
print("saved", name, out.shape, "finite:", np.isfinite(out).all())