"""
Generate image and text embeddings for THINGS-EEG2 dataset.
"""

import logging
from abc import ABC, abstractmethod
from collections.abc import Callable
from os import getenv
from pathlib import Path
from typing import Any

import torch
from diffusers import AutoPipelineForText2Image
from PIL import Image
from safetensors.torch import save_file
from torch import Tensor
from tqdm import tqdm
from transformers import (
    AutoImageProcessor,
    BatchFeature,
    CLIPImageProcessor,
    CLIPModel,
    CLIPTokenizer,
    CLIPVisionModelWithProjection,
    Dinov2WithRegistersModel,
    SiglipModel,
)

from things_eeg2_dataset.cli.main import EmbeddingModel, Partition
from things_eeg2_dataset.paths import layout

from .embedding_models import (
    ImageProjModel,
    IPAdapterPlusXL,
    IPAdapterXL,
    Resampler,
)

__all__ = [
    "BaseEmbedder",
    "DinoV2Embedder",
    "IPAdapterEmbedder",
    "OpenAIClipVitL14Embedder",
    "OpenClipViTH14Embedder",
    "Siglip2Embedder",
    "SiglipEmbedder",
    "build_embedder",
]
logger = logging.getLogger(__name__)


class BaseEmbedder(ABC):
    """Abstract base class for embedding generation."""

    def __init__(
        self,
        project_dir: Path,
        overwrite: bool = False,
        dry_run: bool = False,
        device: str = "cuda:0",
    ) -> None:
        self.images_dir = layout.get_images_dir(project_dir)
        self.train_images_path = layout.get_training_images_dir(project_dir)
        self.test_images_path = layout.get_test_images_dir(project_dir)
        self.embeds_dir = layout.get_embeddings_dir(project_dir)
        self.embeds_dir.mkdir(parents=True, exist_ok=True)

        self.overwrite: bool = overwrite
        self.dry_run: bool = dry_run
        self.device: str = device
        self.model_type: str = "base"

    @staticmethod
    def get_texts(image_dir: Path | str) -> list[str]:
        """Extracts text descriptions from category directory names."""
        image_dir = Path(image_dir)
        text: list[str] = []
        category_dirs: list[Path] = sorted(image_dir.iterdir())
        for category_dir in category_dirs:
            cat_name: str = " ".join(category_dir.name.split("_")[1:])
            text.append(cat_name)
        return text

    @staticmethod
    def get_images(image_dir: Path | str) -> list[Path]:
        """Gathers all image paths from a directory."""
        image_dir = Path(image_dir)
        img_paths: list[Path] = []
        category_dirs: list[Path] = sorted(image_dir.iterdir())
        for category_dir in category_dirs:
            for image_file in sorted(category_dir.iterdir()):
                if image_file.suffix.lower() in (".png", ".jpg", ".jpeg"):
                    im_path: Path = image_file
                    if not im_path.exists():
                        raise FileNotFoundError(f"Image path {im_path} does not exist")
                    img_paths.append(im_path)
        return img_paths

    def store_embeddings(  # noqa: PLR0913
        self,
        image_dir: Path | str,
        embeds_path: Path | str,
        image2embed: Callable[[list[Image.Image]], Tensor],
        txt2embed: Callable[[list[str]], Tensor],
        embed_dim: tuple[int, ...],
        desc: str = "Processing images",
    ) -> None:
        """
        Generates and stores image and text embeddings.

        Args:
            image_dir (Path | str): Path to the directory of images.
            embeds_path (Path | str): Path to save the embeddings.
            image2embed (Callable): Function to convert images to embeddings.
            txt2embed (Callable): Function to convert text to embeddings.
            embed_dim (tuple[int, ...]): The dimensions of the image embeddings.
            desc (str): Description for the progress bar.
        """
        image_dir = Path(image_dir)
        embeds_path = Path(embeds_path)
        Path.mkdir(Path(embeds_path).parent, parents=True, exist_ok=True)
        category_dirs: list[Path] = sorted(image_dir.iterdir())
        num_categories: int = len(category_dirs)
        num_images: int = len(list(category_dirs[0].iterdir()))
        img_embeds: Tensor = torch.zeros(
            num_categories * num_images, *embed_dim, dtype=torch.float16
        )
        text: list[str] = self.get_texts(image_dir)
        img_paths: list[Path] = self.get_images(image_dir)
        if len(img_paths) != num_categories * num_images:
            raise ValueError(
                f"Image paths do not match expected count: {len(img_paths)} but expected {num_categories * num_images}"
            )

        with torch.inference_mode():
            text_embeds: Tensor = txt2embed(text).cpu()
            batch_size = 40
            num_batches = (len(img_paths) + batch_size - 1) // batch_size

            # Create progress bar with proper configuration
            pbar = tqdm(
                total=num_batches,
                desc=desc,
                unit="batch",
                leave=False,
                ncols=100,
                position=0,
            )

            for i in range(0, len(img_paths), batch_size):
                batch_paths: list[Path] = img_paths[i : i + batch_size]
                images: list[Image.Image] = [
                    Image.open(path).convert("RGB") for path in batch_paths
                ]
                img_embeds[i : i + batch_size] = image2embed(images).cpu()
                pbar.update(1)

            pbar.close()

        logger.debug(f"Image embeddings shape: {img_embeds.shape}")
        logger.debug(f"Text embeddings shape: {text_embeds.shape}")

        if self.dry_run:
            logger.info(f"Dry run enabled, not saving embeddings to {embeds_path}")
            return

        embeds_path = embeds_path.with_suffix(".safetensors")

        save_file(
            {
                "img_features": img_embeds,
                "text_features": text_embeds,
            },
            embeds_path,
        )
        logger.info(f"Saved embeddings to {embeds_path}")

    @abstractmethod
    def generate_and_store_embeddings(self, dry_run: bool = False) -> None:
        """An abstract method for generating and storing all embeddings."""
        raise NotImplementedError

    def _store_embeddings_if_needed(  # noqa: PLR0913
        self,
        image_dir: Path | str,
        embeds_path: Path,
        image2embed: Callable[[list[Image.Image]], Tensor],
        txt2embed: Callable[[list[str]], Tensor],
        embed_dim: tuple[int, ...],
        desc: str = "Processing images",
    ) -> None:
        if embeds_path.exists() and not self.overwrite:
            logger.info(f"Embeddings already exist at {embeds_path}, skipping.")
            return
        self.store_embeddings(
            image_dir, embeds_path, image2embed, txt2embed, embed_dim, desc=desc
        )


class OpenClipViTH14Embedder(BaseEmbedder):
    """A class to generate embeddings using OpenCLIP ViT-H-14 model."""

    def __init__(
        self,
        project_dir: Path,
        overwrite: bool = False,
        dry_run: bool = False,
        device: str = "cuda:0",
    ) -> None:
        super().__init__(project_dir, overwrite, dry_run, device)
        self.model_type: str = "ViT-H-14"

        image_encoder: CLIPVisionModelWithProjection = (
            CLIPVisionModelWithProjection.from_pretrained(
                "h94/IP-Adapter",
                subfolder="models/image_encoder",
                torch_dtype=torch.float16,
            ).to(self.device)
        )

        from transformers import CLIPTextModelWithProjection  # noqa: PLC0415

        self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(
            "laion/CLIP-ViT-H-14-laion2B-s32B-b79K"
        )
        self.text_encoder = CLIPTextModelWithProjection.from_pretrained(
            "laion/CLIP-ViT-H-14-laion2B-s32B-b79K",
            torch_dtype=torch.float16,
        ).to(self.device)

        self.feature_extractor: CLIPImageProcessor = CLIPImageProcessor.from_pretrained(
            "laion/CLIP-ViT-H-14-laion2B-s32B-b79K"
        )
        self.image_encoder: CLIPVisionModelWithProjection = image_encoder
        self.project_dir = project_dir

        # Clear cache after loading
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    def transform_clip_vis(self, images: list[Image.Image]) -> Tensor:
        """Generates pooled image embeddings."""
        images_processed: BatchFeature = self.feature_extractor(images)
        pixel_values: list[Any] = images_processed.pixel_values
        images_tensor: Tensor = torch.tensor(pixel_values, device=self.device)
        return self.image_encoder(images_tensor).image_embeds.half()

    def transform_clip_text(self, texts: list[str]) -> Tensor:
        """Generates pooled text embeddings."""
        inputs = self.tokenizer(
            texts,
            padding="max_length",
            max_length=self.tokenizer.model_max_length,
            return_tensors="pt",
        )
        return self.text_encoder(inputs.input_ids.to(self.device)).text_embeds.half()

    def transform_clip_vis_full(self, images: list[Image.Image]) -> Tensor:
        """Generates full sequence image embeddings."""
        images_processed = self.feature_extractor(images)
        pixel_values = images_processed.pixel_values
        images_tensor: Tensor = torch.tensor(pixel_values, device=self.device)
        return self.image_encoder(images_tensor).last_hidden_state.half()

    def transform_clip_text_full(self, texts: list[str]) -> Tensor:
        """Generates full sequence text embeddings."""
        inputs = self.tokenizer(
            texts,
            padding="max_length",
            max_length=self.tokenizer.model_max_length,
            return_tensors="pt",
        )
        return self.text_encoder(
            inputs.input_ids.to(self.device)
        ).last_hidden_state.half()

    def generate_and_store_embeddings(self, dry_run: bool = False) -> None:
        """Generates and stores all required embeddings for the model."""
        # Standard embedding
        train_embeds_path = layout.get_embedding_file(
            self.project_dir,
            self.model_type,
            Partition.TRAINING,
            full=False,
            variant="pooled",
        )
        test_embeds_path = layout.get_embedding_file(
            self.project_dir,
            self.model_type,
            Partition.TEST,
            full=False,
            variant="pooled",
        )
        self._store_embeddings_if_needed(
            self.test_images_path,
            test_embeds_path,
            self.transform_clip_vis,
            self.transform_clip_text,
            embed_dim=(1024,),
            desc="[OpenCLIP ViT-H-14] Test pooled embeddings",
        )
        self._store_embeddings_if_needed(
            self.train_images_path,
            train_embeds_path,
            self.transform_clip_vis,
            self.transform_clip_text,
            embed_dim=(1024,),
            desc="[OpenCLIP ViT-H-14] Train pooled embeddings",
        )

        # Full token embedding
        train_embeds_path = layout.get_embedding_file(
            self.project_dir,
            self.model_type,
            Partition.TRAINING,
            full=True,
            variant="full",
        )
        test_embeds_path = layout.get_embedding_file(
            self.project_dir, self.model_type, Partition.TEST, full=True, variant="full"
        )
        self._store_embeddings_if_needed(
            self.test_images_path,
            test_embeds_path,
            self.transform_clip_vis_full,
            self.transform_clip_text_full,
            embed_dim=(257, 1280),
            desc="[OpenCLIP ViT-H-14] Test full embeddings",
        )
        self._store_embeddings_if_needed(
            self.train_images_path,
            train_embeds_path,
            self.transform_clip_vis_full,
            self.transform_clip_text_full,
            embed_dim=(257, 1280),
            desc="[OpenCLIP ViT-H-14] Train full embeddings",
        )


class OpenAIClipVitL14Embedder(BaseEmbedder):
    """A class to generate embeddings using OpenAI's CLIP ViT-L-14 model."""

    def __init__(
        self,
        project_dir: Path,
        overwrite: bool = False,
        dry_run: bool = False,
        device: str = "cuda:0",
    ) -> None:
        super().__init__(project_dir, overwrite, dry_run, device)
        self.model_type = "openai-clip-vit-l-14"
        self.processor: CLIPImageProcessor = CLIPImageProcessor.from_pretrained(
            "openai/clip-vit-large-patch14"
        )
        self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(
            "openai/clip-vit-large-patch14"
        )
        self.model: CLIPModel = CLIPModel.from_pretrained(
            "openai/clip-vit-large-patch14"
        )
        self.model.to(device)
        self.project_dir = project_dir

    def encode_text_pooled(self, text: list[str]) -> Tensor:
        """Generates pooled text embeddings."""
        batch_encoding: dict[str, Any] = self.tokenizer(
            text,
            truncation=True,
            max_length=77,
            return_length=True,
            return_overflowing_tokens=False,
            padding="max_length",
            return_tensors="pt",
        )
        tokens: Tensor = batch_encoding["input_ids"].to(self.device)
        return self.model.get_text_features(input_ids=tokens)

    def encode_vision_pooled(self, images: list[Image.Image]) -> Tensor:
        """Generates pooled image embeddings."""
        inputs = self.processor(images=images, return_tensors="pt")
        pixels: Tensor = inputs["pixel_values"].to(self.device)
        return self.model.get_image_features(pixel_values=pixels)

    def encode_text(self, text: list[str]) -> Tensor:
        """Generates full sequence text embeddings."""
        batch_encoding: dict[str, Any] = self.tokenizer(
            text,
            truncation=True,
            max_length=77,
            return_length=True,
            return_overflowing_tokens=False,
            padding="max_length",
            return_tensors="pt",
        )
        tokens: Tensor = batch_encoding["input_ids"].to(self.device)
        outputs: Any = self.model.text_model(input_ids=tokens)
        z: Tensor = self.model.text_projection(outputs.last_hidden_state)
        z_pooled: Tensor = self.model.text_projection(outputs.pooler_output)
        z = z / torch.norm(z_pooled.unsqueeze(1), dim=-1, keepdim=True)
        return z

    def encode_vision_noproj(self, images: list[Image.Image]) -> Tensor:
        """Helper for full sequence image embeddings without projection."""
        inputs = self.processor(images=images, return_tensors="pt")
        pixels: Tensor = inputs["pixel_values"].to(self.device)
        outputs: Any = self.model.vision_model(pixel_values=pixels)
        return outputs.last_hidden_state

    def encode_vision(self, images: list[Image.Image]) -> Tensor:
        """Generates full sequence image embeddings."""
        z: Tensor = self.encode_vision_noproj(images)
        z = self.model.vision_model.post_layernorm(z)
        z = self.model.visual_projection(z)
        z_pooled: Tensor = z[:, 0:1]
        z = z / torch.norm(z_pooled, dim=-1, keepdim=True)
        return z

    def generate_and_store_embeddings(self, dry_run: bool = False) -> None:
        """Generates and stores all required embeddings for the model."""
        # Pooled embedding
        train_embeds_path = layout.get_embedding_file(
            self.project_dir,
            self.model_type,
            Partition.TRAINING,
            full=False,
            variant="pooled",
        )
        test_embeds_path = layout.get_embedding_file(
            self.project_dir,
            self.model_type,
            Partition.TEST,
            full=False,
            variant="pooled",
        )
        self._store_embeddings_if_needed(
            self.test_images_path,
            test_embeds_path,
            self.encode_vision_pooled,
            self.encode_text_pooled,
            embed_dim=(768,),
            desc="[OpenAI CLIP ViT-L-14] Test pooled embeddings",
        )
        self._store_embeddings_if_needed(
            self.train_images_path,
            train_embeds_path,
            self.encode_vision_pooled,
            self.encode_text_pooled,
            embed_dim=(768,),
            desc="[OpenAI CLIP ViT-L-14] Train pooled embeddings",
        )

        # Full token embedding
        train_embeds_path = layout.get_embedding_file(
            self.project_dir,
            self.model_type,
            Partition.TRAINING,
            full=True,
            variant="full",
        )
        test_embeds_path = layout.get_embedding_file(
            self.project_dir, self.model_type, Partition.TEST, full=True, variant="full"
        )
        self._store_embeddings_if_needed(
            self.test_images_path,
            test_embeds_path,
            self.encode_vision,
            self.encode_text,
            embed_dim=(257, 768),
            desc="[OpenAI CLIP ViT-L-14] Test full embeddings",
        )
        self._store_embeddings_if_needed(
            self.train_images_path,
            train_embeds_path,
            self.encode_vision,
            self.encode_text,
            embed_dim=(257, 768),
            desc="[OpenAI CLIP ViT-L-14] Train full embeddings",
        )


class DinoV2Embedder(BaseEmbedder):
    """A class to generate embeddings using DINOv2 with registers."""

    def __init__(
        self,
        project_dir: Path,
        overwrite: bool = False,
        dry_run: bool = False,
        device: str = "cuda:0",
    ) -> None:
        super().__init__(project_dir, overwrite, dry_run, device)
        self.model_type = "dinov2-reg"
        self.image_processor: AutoImageProcessor = AutoImageProcessor.from_pretrained(
            "facebook/dinov2-with-registers-base", use_fast=True
        )
        self.dinov2_model: Dinov2WithRegistersModel = (
            Dinov2WithRegistersModel.from_pretrained(
                "facebook/dinov2-with-registers-base"
            )
        )
        self.dinov2_model.to(device)
        self.project_dir = project_dir

    def transform_dinov2_vis_full(self, images: list[Image.Image]) -> Tensor:
        """Generates full sequence image embeddings."""
        images_processed = self.image_processor(images, return_tensors="pt")
        pixel_values: Tensor = images_processed.pixel_values
        images_tensor: Tensor = torch.tensor(pixel_values, device=self.device)
        tokens: Tensor = self.dinov2_model(images_tensor).last_hidden_state
        return tokens.half()

    def transform_dinov2_vis_cls(self, images: list[Image.Image]) -> Tensor:
        """Generates CLS token image embeddings."""
        tokens: Tensor = self.transform_dinov2_vis_full(images)
        return tokens[:, 0, :]

    def transform_dinov2_vis_reg(self, images: list[Image.Image]) -> Tensor:
        """Generates CLS + register tokens image embeddings."""
        tokens: Tensor = self.transform_dinov2_vis_full(images)
        return tokens[:, 0:5, :]

    def transform_dinov2_text(self, texts: list[str]) -> Tensor:
        """Returns a dummy tensor as DINOv2 is vision-only."""
        logger.warning(
            "WARNING: DINOv2 cannot be used for text, returning dummy tensor"
        )
        return torch.tensor(0.0, device=self.device).half()

    def generate_and_store_embeddings(self, dry_run: bool = False) -> None:
        """Generates and stores all required embeddings for the model."""
        # CLS token embedding
        train_embeds_path = layout.get_embedding_file(
            self.project_dir,
            self.model_type,
            Partition.TRAINING,
            full=False,
            variant="cls",
        )
        test_embeds_path = layout.get_embedding_file(
            self.project_dir, self.model_type, Partition.TEST, full=False, variant="cls"
        )
        self._store_embeddings_if_needed(
            self.test_images_path,
            test_embeds_path,
            self.transform_dinov2_vis_cls,
            self.transform_dinov2_text,
            embed_dim=(768,),
            desc="[DINOv2] Test CLS embeddings",
        )
        self._store_embeddings_if_needed(
            self.train_images_path,
            train_embeds_path,
            self.transform_dinov2_vis_cls,
            self.transform_dinov2_text,
            embed_dim=(768,),
            desc="[DINOv2] Train CLS embeddings",
        )

        # CLS + Register tokens embedding
        train_embeds_path = layout.get_embedding_file(
            self.project_dir,
            self.model_type,
            Partition.TRAINING,
            full=False,
            variant="cls+register",
        )
        test_embeds_path = layout.get_embedding_file(
            self.project_dir,
            self.model_type,
            Partition.TEST,
            full=False,
            variant="cls+register",
        )
        self._store_embeddings_if_needed(
            self.test_images_path,
            test_embeds_path,
            self.transform_dinov2_vis_reg,
            self.transform_dinov2_text,
            embed_dim=(5, 768),
            desc="[DINOv2] Test CLS+Register embeddings",
        )
        self._store_embeddings_if_needed(
            self.train_images_path,
            train_embeds_path,
            self.transform_dinov2_vis_reg,
            self.transform_dinov2_text,
            embed_dim=(5, 768),
            desc="[DINOv2] Train CLS+Register embeddings",
        )

        # All tokens embedding
        train_embeds_path = layout.get_embedding_file(
            self.project_dir,
            self.model_type,
            Partition.TRAINING,
            full=True,
            variant="all",
        )
        test_embeds_path = layout.get_embedding_file(
            self.project_dir, self.model_type, Partition.TEST, full=True, variant="all"
        )
        self._store_embeddings_if_needed(
            self.test_images_path,
            test_embeds_path,
            self.transform_dinov2_vis_full,
            self.transform_dinov2_text,
            embed_dim=(261, 768),
            desc="[DINOv2] Test all tokens embeddings",
        )
        self._store_embeddings_if_needed(
            self.train_images_path,
            train_embeds_path,
            self.transform_dinov2_vis_full,
            self.transform_dinov2_text,
            embed_dim=(261, 768),
            desc="[DINOv2] Train all tokens embeddings",
        )


class IPAdapterEmbedder(BaseEmbedder):
    """A class to generate embeddings using IP-Adapter with ViT-H-14."""

    def __init__(
        self,
        project_dir: Path,
        overwrite: bool = False,
        dry_run: bool = False,
        device: str = "cuda:0",
    ) -> None:
        super().__init__(project_dir, overwrite, dry_run, device)
        self.model_type = "ip-adapter-plus-vit-h-14"
        image_encoder: CLIPVisionModelWithProjection = (
            CLIPVisionModelWithProjection.from_pretrained(
                "h94/IP-Adapter",
                subfolder="models/image_encoder",
                torch_dtype=torch.float16,
            )
        )
        pipeline: Any = AutoPipelineForText2Image.from_pretrained(
            "stabilityai/stable-diffusion-xl-base-1.0",
            image_encoder=image_encoder,
            torch_dtype=torch.float16,
        ).to(device)
        pipeline.load_ip_adapter(
            "h94/IP-Adapter",
            subfolder="sdxl_models",
            weight_name="ip-adapter-plus_sdxl_vit-h.safetensors",
        )

        self.ip_adapter_proj: Resampler = self._load_ipa_proj_model(
            plus_model=True, device=device
        )
        self.ip_adapter_proj.to(device).eval()
        self.feature_extractor: CLIPImageProcessor = pipeline.feature_extractor
        self.image_encoder = pipeline.image_encoder
        self.project_dir = project_dir

    def _load_ipa_proj_model(
        self, plus_model: bool = True, device: str = "cuda:0"
    ) -> ImageProjModel | Resampler:
        base_model_path = "stabilityai/stable-diffusion-xl-base-1.0"
        image_encoder_path = "models/image_encoder"

        if plus_model:
            num_tokens = 16
            fname = "ip-adapter-plus_sdxl_vit-h.safetensors"
        else:
            num_tokens = 4
            fname = "ip-adapter_sdxl_vit-h.safetensors"

        # Load cached weights from huggingface hub
        hf_home = Path(getenv("HF_HOME", str(Path.home() / ".cache/huggingface")))
        weights_path = hf_home / "hub" / "models--h94--IP-Adapter" / "snapshots"
        weights_path = weights_path / "018e402774aeeddd60609b4ecdb7e298259dc729"
        weights_path = weights_path / "sdxl_models" / fname

        # Create the diffusion pipeline
        pipe = AutoPipelineForText2Image.from_pretrained(
            base_model_path, torch_dtype=torch.float16
        )

        ip_model: IPAdapterPlusXL | IPAdapterXL

        if plus_model:
            ip_model = IPAdapterPlusXL(
                pipe,
                image_encoder_path,
                str(weights_path),
                device,
                num_tokens=num_tokens,
            )
        else:
            ip_model = IPAdapterXL(
                pipe,
                image_encoder_path,
                str(weights_path),
                device,
                num_tokens=num_tokens,
            )

        # Extract image projection model from IP-Adapter
        image_proj_model = ip_model.image_proj_model
        image_proj_model.requires_grad_(False)
        image_proj_model.half()

        del pipe
        torch.cuda.empty_cache()

        return image_proj_model

    def transform_ipa_vis(self, images: list[Image.Image]) -> Tensor:
        """Generates IP-Adapter image embeddings."""
        images_processed: BatchFeature = self.feature_extractor(
            images, return_tensors="pt"
        )
        pixel_values: Tensor = images_processed.pixel_values
        images_tensor: Tensor = torch.tensor(
            pixel_values, device=self.device, dtype=torch.float16
        )
        clip_embeds: Tensor = self.image_encoder(images_tensor).last_hidden_state
        ipa_embeds: Tensor = self.ip_adapter_proj(clip_embeds)
        return ipa_embeds.half()

    def transform_dummy_text(self, texts: list[str]) -> Tensor:
        """Returns a dummy tensor as IP-Adapter is vision-only."""
        logger.warning(
            "WARNING: IP-Adapter is for images, returning dummy tensor for text"
        )
        return torch.tensor(0.0, device=self.device).half()

    def generate_and_store_embeddings(self, dry_run: bool = False) -> None:
        """Generates and stores all required embeddings for the model."""
        train_embeds_path: Path = (
            self.embeds_dir / f"{self.model_type}_features_train.pt"
        )
        test_embeds_path: Path = self.embeds_dir / f"{self.model_type}_features_test.pt"
        self._store_embeddings_if_needed(
            self.test_images_path,
            test_embeds_path,
            self.transform_ipa_vis,
            self.transform_dummy_text,
            embed_dim=(16, 2048),
            desc="[IP-Adapter] Test embeddings",
        )
        self._store_embeddings_if_needed(
            self.train_images_path,
            train_embeds_path,
            self.transform_ipa_vis,
            self.transform_dummy_text,
            embed_dim=(16, 2048),
            desc="[IP-Adapter] Train embeddings",
        )


class SiglipEmbedder(BaseEmbedder):
    """A class to generate embeddings using Siglip model."""

    def __init__(
        self,
        project_dir: Path,
        overwrite: bool = False,
        dry_run: bool = False,
        device: str = "cuda:0",
    ) -> None:
        super().__init__(project_dir, overwrite, dry_run, device)
        self.model_type = "siglip-base-patch16-224"

        # Loading of the pretrained SigLIP model
        self.model: SiglipModel = SiglipModel.from_pretrained(
            "google/siglip-base-patch16-224"
        )
        self.model.to(device)

        # Load processor and tokenizer from the pretrained model

        from transformers import AutoProcessor  # noqa: PLC0415

        self.processor = AutoProcessor.from_pretrained("google/siglip-base-patch16-224")

        self.project_dir = project_dir

    def transform_siglip_vis(self, images: list[Image.Image]) -> Tensor:
        """Generates pooled image embeddings."""
        inputs = self.processor(images=images, return_tensors="pt")
        pixels: Tensor = inputs["pixel_values"].to(self.device)
        return self.model.get_image_features(pixel_values=pixels)

    def transform_siglip_text(self, texts: list[str]) -> Tensor:
        """Generates text embeddings."""
        inputs = self.processor(
            texts,
            padding="max_length",
            max_length=64,
            return_tensors="pt",
        )
        tokens: Tensor = inputs["input_ids"].to(self.device)  # noqa: F841
        return self.model.get_text_features(
            input_ids=inputs["input_ids"].to(self.device)
        )

    def generate_and_store_embeddings(self, dry_run: bool = False) -> None:
        """Generates and stores all required embeddings for the model."""
        # Pooled embedding
        train_embeds_path = layout.get_embedding_file(
            self.project_dir,
            self.model_type,
            Partition.TRAINING,
            full=False,
            variant="pooled",
        )
        test_embeds_path = layout.get_embedding_file(
            self.project_dir,
            self.model_type,
            Partition.TEST,
            full=False,
            variant="pooled",
        )
        self._store_embeddings_if_needed(
            self.train_images_path,
            train_embeds_path,
            self.transform_siglip_vis,
            self.transform_siglip_text,
            embed_dim=(768,),
            desc="[SigLIP] Train pooled embeddings",
        )

        self._store_embeddings_if_needed(
            self.test_images_path,
            test_embeds_path,
            self.transform_siglip_vis,
            self.transform_siglip_text,
            embed_dim=(768,),
            desc="[SigLIP] Test pooled embeddings",
        )


class Siglip2Embedder(BaseEmbedder):
    """A class to generate embeddings using Siglip2 model."""

    def __init__(
        self,
        project_dir: Path,
        overwrite: bool = False,
        dry_run: bool = False,
        device: str = "cuda:0",
    ) -> None:
        super().__init__(project_dir, overwrite, dry_run, device)
        self.model_type = "siglip2-base-patch16-224"

        # Load the pretrained SigLIP2 model
        from transformers import AutoModel, AutoProcessor  # noqa: PLC0415

        self.model = AutoModel.from_pretrained(
            "google/siglip2-base-patch16-224", torch_dtype=torch.float16
        )
        self.model.to(device)
        self.model.eval()  # Set to evaluation mode

        # Load processor from the pretrained model
        self.processor = AutoProcessor.from_pretrained(
            "google/siglip2-base-patch16-224"
        )

        self.project_dir = project_dir

        # Clear cache after loading
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    def transform_siglip2_vis(self, images: list[Image.Image]) -> Tensor:
        """Generates pooled image embeddings."""
        inputs = self.processor(images=images, return_tensors="pt")
        pixels: Tensor = inputs["pixel_values"].to(self.device, dtype=torch.float16)

        vision_outputs = self.model.vision_model(pixel_values=pixels)

        return vision_outputs.pooler_output.half()

    def transform_siglip2_text(self, texts: list[str]) -> Tensor:
        """Generates text embeddings."""
        inputs = self.processor(
            text=texts,
            padding="max_length",
            max_length=64,
            return_tensors="pt",
        )

        input_ids = inputs["input_ids"].to(self.device)
        attention_mask = inputs.get("attention_mask")
        if attention_mask is not None:
            attention_mask = attention_mask.to(self.device)

        text_outputs = self.model.text_model(
            input_ids=input_ids, attention_mask=attention_mask
        )

        return text_outputs.pooler_output.half()

    def generate_and_store_embeddings(self, dry_run: bool = False) -> None:
        """Generates and stores all required embeddings for the model."""
        # Pooled embedding
        train_embeds_path = layout.get_embedding_file(
            self.project_dir,
            self.model_type,
            Partition.TRAINING,
            full=False,
            variant="pooled",
        )
        test_embeds_path = layout.get_embedding_file(
            self.project_dir,
            self.model_type,
            Partition.TEST,
            full=False,
            variant="pooled",
        )

        logger.info("Generating SigLIP2 embeddings for training set...")
        self._store_embeddings_if_needed(
            self.train_images_path,
            train_embeds_path,
            self.transform_siglip2_vis,
            self.transform_siglip2_text,
            embed_dim=(768,),
            desc="[SigLIP2] Train pooled embeddings",
        )

        logger.info("Generating SigLIP2 embeddings for test set...")
        self._store_embeddings_if_needed(
            self.test_images_path,
            test_embeds_path,
            self.transform_siglip2_vis,
            self.transform_siglip2_text,
            embed_dim=(768,),
            desc="[SigLIP2] Test pooled embeddings",
        )


EMBEDDER_DICT = {
    EmbeddingModel.OPEN_CLIP_VIT_H_14: OpenClipViTH14Embedder,
    EmbeddingModel.OPENAI_CLIP_VIT_L_14: OpenAIClipVitL14Embedder,
    EmbeddingModel.DINO_V2: DinoV2Embedder,
    EmbeddingModel.IP_ADAPTER: IPAdapterEmbedder,
    EmbeddingModel.SIGLIP: SiglipEmbedder,
    EmbeddingModel.SIGLIP2: Siglip2Embedder,
}


def build_embedder(
    model_type: EmbeddingModel,
    project_dir: Path,
    overwrite: bool = False,
    dry_run: bool = False,
    device: str = "cuda:0",
) -> BaseEmbedder:
    """Factory function to build the appropriate embedder based on model type."""
    embedder_class = EMBEDDER_DICT[model_type]
    return embedder_class(project_dir, overwrite, dry_run, device)  # type: ignore
