from __future__ import annotations
import torch
import open_clip
from contextlib import nullcontext

from src.models.utils import l2norm_rows

class CLIPMultiLabel:
    def __init__(self, head_path, categories):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.torch_dtype = torch.float16 if self.device == "cuda" else torch.float32
        self.categories = list(categories)
        self.model, _, self.preprocess = open_clip.create_model_and_transforms(
            "ViT-L-14", pretrained="openai", device=self.device
        )
        self.model.eval().requires_grad_(False)

        ckpt = torch.load(head_path, map_location=self.device, weights_only=True)
        state = ckpt.get("state_dict", ckpt)

        w = state["head.weight"].to(self.device).float()
        b = state["head.bias"].to(self.device).float()
        w = w.t()

        self.w, self.b = w, b

        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)
        return l2norm_rows(f.float())

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

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