"""SigLIP2-B/16 NaFlex vision tower (frozen), fixed 384x384 -> 24x24 = 576 patches x 768.
encode(images) -> fp16 numpy [N,576,768]: last_hidden_state of the vision tower = AFTER the final post-layernorm,
BEFORE the attention-pooling head (the head is not used).  Pixel prep: resize to 384x384 (bicubic, aspect NOT kept),
rescale 1/255, normalise mean=std=0.5, patchify row-major (16x16x3 -> 768, order (ph,pw,c))."""
import os, numpy as np, torch
from PIL import Image
HERE = os.path.dirname(os.path.abspath(__file__))
# Local copy of the encoder if SIGLIP2_PATH is set; otherwise the public Hub repo (auto-downloaded).
MODEL_DIR = os.environ.get("SIGLIP2_PATH") or os.path.join(HERE, "siglip2")
if not os.path.isdir(MODEL_DIR): MODEL_DIR = "google/siglip2-base-patch16-naflex"
S, P = 384, 16; G = S // P  # 24

_model = None
def load(device="cpu", dtype=torch.float16):
    global _model
    from transformers import Siglip2VisionModel
    m = Siglip2VisionModel.from_pretrained(MODEL_DIR, dtype=dtype).to(device).eval()
    for p in m.parameters(): p.requires_grad_(False)
    _model = (m, device, dtype); return _model

def prep(img):
    if not isinstance(img, Image.Image): img = Image.open(img)
    img = img.convert("RGB").resize((S, S), Image.BICUBIC)
    a = (np.asarray(img, dtype=np.float32) / 255.0 - 0.5) / 0.5            # [384,384,3]
    a = a.reshape(G, P, G, P, 3).transpose(0, 2, 1, 3, 4).reshape(G * G, P * P * 3)  # [576,768]
    return a

@torch.no_grad()
def encode(images, device="cpu", batch=8):
    """images: list of PIL images / paths -> np.float16 [N,576,768]"""
    global _model
    if _model is None or _model[1] != device: load(device)
    m, dev, dt = _model; out = []
    for i in range(0, len(images), batch):
        x = torch.from_numpy(np.stack([prep(im) for im in images[i:i + batch]])).to(dev, dt)
        b = x.shape[0]
        mask = torch.ones(b, G * G, dtype=torch.long, device=dev)
        shapes = torch.tensor([[G, G]] * b, dtype=torch.long, device=dev)
        h = m(pixel_values=x, pixel_attention_mask=mask, spatial_shapes=shapes).last_hidden_state
        out.append(h.float().cpu().numpy().astype(np.float16))
    return np.concatenate(out)