"""Crop EMBEDDER for the concept bag — model-agnostic (CLIP/SigLIP-family).

The server trains its per-concept heads in the embedding space of whatever model
its `embedder_model_version` names; a crop must be embedded with the SAME model
or its vector lands in a different coordinate system and every head misfires. So
the model identity (HF name + version) is ANNOUNCED BY THE SERVER in the lease —
nothing here is hardcoded to SigLIP. Whatever name the server sends is loaded via
transformers `get_image_features` (the CLIP/SigLIP-family image-tower call); a
non-CLIP backbone (e.g. a DINO encoder) would need its own pooling adapter.

torch on CUDA, fp16 by default to keep VRAM low on a shared desktop GPU — the
tiny fp16-vs-fp32 difference is negligible for the linear heads (cosine ~0.999).
A single inference lock serializes the forward pass: the pipeline is I/O-bound,
so the GPU isn't the bottleneck, and one model shared across worker threads is
safest behind a lock.
"""
import threading

import numpy as np
from PIL import Image


class CropEmbedder:
    def __init__(self, model_name: str, dtype: str = "float16"):
        self._name = model_name
        self._dtype_name = dtype
        self._model = None
        self._processor = None
        self._torch = None
        self._device = None
        self._dt = None
        self._load_lock = threading.Lock()
        self._infer_lock = threading.Lock()

    @property
    def model_name(self) -> str:
        return self._name

    def load(self) -> None:
        if self._model is not None:
            return
        with self._load_lock:
            if self._model is not None:
                return
            import torch
            from transformers import AutoImageProcessor, AutoModel

            self._torch = torch
            self._device = "cuda" if torch.cuda.is_available() else "cpu"
            dt = getattr(torch, self._dtype_name, torch.float16)
            if self._device == "cpu":
                dt = torch.float32   # fp16 matmul is unsupported/slow on CPU
            self._dt = dt
            self._processor = AutoImageProcessor.from_pretrained(self._name)
            model = AutoModel.from_pretrained(self._name, torch_dtype=dt)
            model.eval().to(self._device)
            self._model = model

    def embed(self, image: Image.Image) -> list[float]:
        """A crop → its embedding as a plain float list, ready to POST."""
        self.load()
        torch = self._torch
        enc = self._processor(images=image, return_tensors="pt")
        pixel_values = enc["pixel_values"].to(self._device, self._dt)
        with self._infer_lock, torch.no_grad():
            out = self._model.get_image_features(pixel_values=pixel_values)
        pooled = out.pooler_output if hasattr(out, "pooler_output") else out
        vec = pooled[0].float().cpu().numpy().astype(np.float32).reshape(-1)
        return vec.tolist()