from __future__ import annotations
import numpy as np
import torch
import open_clip
from contextlib import nullcontext

from src.models.utils import l2norm_rows

class CLIPLinearProbe:
    def __init__(self, head_path):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.torch_dtype = torch.float16 if self.device == "cuda" else torch.float32

        self.model, _, self.preprocess = open_clip.create_model_and_transforms(
            "ViT-L-14", pretrained="openai", device=self.device
        )
        self.model.eval().requires_grad_(False)
        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:
        x = torch.stack([self.preprocess(im.convert("RGB")) for im in pil_list], 0)
        x = x.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.encode_image(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)