"""Frozen text-tower embeddings of the templated tile descriptions.
Usage: python 08_extract_text.py {plip|biomedclip|imagebind|conchclip} [--set <name> | --pathgen | --prompts]
Output: cache/text_<enc>.npy (N, 4 paraphrases, D) float16.
  plip       : vinid/plip CLIP text tower (77 tokens), projected -> 512
  biomedclip : microsoft/BiomedCLIP PubMedBERT text tower (256 tokens), projected -> 512
  imagebind  : ImageBind-Huge text tower (OpenCLIP ViT-H text, 77 tokens) -> 1024
  conchclip  : MahmoodLab/CONCH text tower (128 tokens), projected -> 512
"""
import json
import os
import sys

import numpy as np
import torch

from common import CACHE, ROOT, 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
PROMPTS = "--prompts" in sys.argv  # encode concept prompts (prompts.json) for zero-shot probing
if PROMPTS:
    pr = json.load(open(os.path.join(ROOT, "prompts.json")))
    flat = [v["text"] for v in pr.values()]
    ids, NV = list(range(len(flat))), 1
elif PATHGEN:
    import pandas as pd
    flat = pd.read_csv(os.path.join(CACHE, f"{SET}_meta.csv")).caption.tolist()
    ids, NV = list(range(len(flat))), 1
else:
    texts = json.load(open(os.path.join(CACHE, "text.json")))
    ids, NV = sorted(int(k) for k in texts), 4
    flat = [texts[str(i)][v] for i in ids for v in range(NV)]

if name == "plip":
    from transformers import CLIPModel, CLIPTokenizer
    m = CLIPModel.from_pretrained("vinid/plip").cuda().eval()
    tok = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")  # PLIP ships no tokenizer; it uses the CLIP B/32 one
    enc = lambda b: m.get_text_features(**{k: v.cuda() for k, v in tok(b, padding=True, truncation=True, max_length=77, return_tensors="pt").items()})
elif name == "biomedclip":
    import open_clip
    m, _ = open_clip.create_model_from_pretrained("hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224")
    tok = open_clip.get_tokenizer("hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224")
    m = m.cuda().eval()
    enc = lambda b: m.encode_text(tok(b, context_length=256).cuda())
elif name == "imagebind":
    model, MT = load_imagebind_vision()
    from imagebind.models.multimodal_preprocessors import SimpleTokenizer
    tk = SimpleTokenizer(bpe_path=os.path.join(ROOT, "ImageBind_src", "imagebind", "bpe", "bpe_simple_vocab_16e6.txt.gz"))
    enc = lambda b: model({MT.TEXT: torch.cat([tk(t).unsqueeze(0) for t in b]).cuda()})[MT.TEXT]
elif name == "conchclip":  # CONCH text tower (CoCa, 128 tokens), projected -> 512 (pairs with the conchclip image tower)
    from conch.open_clip_custom import create_model_from_pretrained, get_tokenizer, tokenize
    m, _ = create_model_from_pretrained("conch_ViT-B-16", "hf_hub:MahmoodLab/conch")
    m = m.cuda().eval()
    tok = get_tokenizer()
    enc = lambda b: m.encode_text(tokenize(tok, b).cuda(), normalize=False)
else:
    raise ValueError(name)

out = []
with torch.no_grad():
    for s in range(0, len(flat), 256):
        out.append(enc(flat[s:s + 256]).float().cpu().numpy())
out = np.concatenate(out).reshape(len(ids), NV, -1).astype(np.float16)
assert ids == list(range(len(ids)))
pre = "prompts_" if PROMPTS else f"{SET}_" if PATHGEN else ""
np.save(os.path.join(CACHE, pre + f"text_{name}.npy"), out)
print("saved", name, out.shape, np.isfinite(out).all())