"""Frozen embeddings of every text set with additional (bio)medical / general text-embedding models.

    python 45_text_feats.py <name> [<name> ...]

Text sets (same row order as the existing BiomedCLIP files):
  alltxt         cache/alltxt_meta.csv caption        (TCGA real PathGen texts, 6480, junk rows kept, filtered later)
  orion_gentext  cache/orion_gentext_meta.csv caption (GPT synthetic captions of the ORION tiles, test only)
  tcga_notext    cache/tcga_notext_meta.csv caption   (synthetic captions of TCGA images without real text, test only)
  dlpfc_gentext  cache/dlpfc_gentext_meta.csv caption (external DLPFC synthetic captions, test only)
Models (pooling / normalisation as on each model card; documents embedded without a query instruction):
  biomedclip   BiomedCLIP text tower (biomedclip_text.py, 256 tokens) - only for the sets that lack it
  medcpt       ncbi/MedCPT-Article-Encoder, [CLS], 512 tokens (PubMed search-log contrastive)
  biolord      FremyCompany/BioLORD-2023, mean pooling, 512 (clinical sentences / concepts, UMLS-based)
  pubmedbert   NeuML/pubmedbert-base-embeddings, mean pooling, 512 (PubMed title-abstract)
  medembed     abhinand/MedEmbed-large-v0.1 (bge-large medical fine-tune), [CLS], 512
  bge          BAAI/bge-large-en-v1.5, [CLS], 512 (general-domain reference)
  qwen3-0.6b / qwen3-4b   Qwen/Qwen3-Embedding-*, last-token pooling, left padding, 1024 tokens
Output cache/<set>_text_<name>.npy (n, 1, D) float16 (raw, not normalised).
"""
import os
import sys
import time

import numpy as np
import pandas as pd
import torch
from transformers import AutoModel, AutoTokenizer

from common import CACHE

dev = "cuda"
SETS = {"alltxt": "alltxt_meta.csv", "orion_gentext": "orion_gentext_meta.csv", "tcga_notext": "tcga_notext_meta.csv",
        "dlpfc_gentext": "dlpfc_gentext_meta.csv"}
HF = {"medcpt": ("ncbi/MedCPT-Article-Encoder", "cls", 512), "biolord": ("FremyCompany/BioLORD-2023", "mean", 512),
      "pubmedbert": ("NeuML/pubmedbert-base-embeddings", "mean", 512),
      "medembed": ("abhinand/MedEmbed-large-v0.1", "cls", 512), "bge": ("BAAI/bge-large-en-v1.5", "cls", 512),
      "qwen3-0.6b": ("Qwen/Qwen3-Embedding-0.6B", "last", 1024), "qwen3-4b": ("Qwen/Qwen3-Embedding-4B", "last", 1024)}


def encoder(name):
    if name == "biomedclip":
        from biomedclip_text import BiomedCLIPText
        m = BiomedCLIPText().to(dev).eval()
        return m, 64
    repo, pool, L = HF[name]
    tk = AutoTokenizer.from_pretrained(repo, padding_side="left" if pool == "last" else "right")
    m = AutoModel.from_pretrained(repo, torch_dtype=torch.bfloat16 if pool == "last" else torch.float32).to(dev).eval()

    def f(b):
        x = tk(b, padding=True, truncation=True, max_length=L, return_tensors="pt").to(dev)
        h = m(**x).last_hidden_state
        if pool == "cls":
            return h[:, 0]
        if pool == "mean":
            a = x["attention_mask"][..., None].to(h.dtype)
            return (h * a).sum(1) / a.sum(1)
        return h[:, -1]  # left padding -> last position is the last real token
    return f, (16 if name == "qwen3-4b" else 64)


for name in sys.argv[1:]:
    t0 = time.time()
    f, bs = encoder(name)
    for st, meta in SETS.items():
        out_f = os.path.join(CACHE, f"{st}_text_{name}.npy")
        if os.path.exists(out_f):
            continue
        txt = pd.read_csv(os.path.join(CACHE, meta)).caption.fillna("").astype(str).tolist()
        order = np.argsort([len(t) for t in txt])  # length-sorted batches, restored below
        E = None
        with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
            for i in range(0, len(txt), bs):
                idx = order[i:i + bs]
                z = f([txt[j] for j in idx]).float().cpu().numpy()
                if E is None:
                    E = np.zeros((len(txt), z.shape[1]), np.float32)
                E[idx] = z
        assert np.isfinite(E).all()
        np.save(out_f, E[:, None].astype(np.float16))
        print(f"[{name}] {st}: {E.shape}  {time.time() - t0:.0f}s", flush=True)
    del f
    torch.cuda.empty_cache()