from __future__ import annotations
import numpy as np
import torch
from contextlib import nullcontext
from transformers import AutoProcessor, SiglipModel

from src.models.utils import l2norm_rows

class SigLIPLinearProbe:
    def __init__(self, head_path):
        self.model_id = "google/siglip-so400m-patch14-384"
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.torch_dtype = torch.float16 if self.device == "cuda" else torch.float32
        self.model = SiglipModel.from_pretrained(self.model_id, dtype=self.torch_dtype).to(self.device)
        self.model.eval().requires_grad_(False)
        self.processor = AutoProcessor.from_pretrained(self.model_id)

        npz = np.load(head_path)
        self.w = torch.from_numpy(npz["w"]).to(self.device).float()
        self.b = torch.from_numpy(npz["b"]).to(self.device).float()

        self.use_amp = False
        if self.device == "cuda":
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.benchmark = True
            self.use_amp = True

    @torch.inference_mode()
    def encode(self, pil_list) -> torch.Tensor:
        imgs = [im.convert("RGB") for im in pil_list]
        enc = self.processor(images=imgs, return_tensors="pt")
        x = enc["pixel_values"].to(self.device, non_blocking=True, memory_format=torch.channels_last)
        ctx = torch.amp.autocast("cuda", dtype=self.torch_dtype) if self.use_amp else nullcontext()
        with ctx:
            f = self.model.get_image_features(pixel_values=x)
        f = f.float()
        return l2norm_rows(f)

    @torch.inference_mode()
    def logits(self, pil_list) -> torch.Tensor:
        f = self.encode(pil_list)
        return (f @ self.w + self.b).squeeze(1)

    @torch.inference_mode()
    def prob(self, pil_list ) -> torch.Tensor:
        z = torch.clamp(self.logits(pil_list), -50, 50)
        return torch.sigmoid(z)