"""
COLIP: Olfaction-Vision-Language Embedding models.

Four variants are available, all hosted at:
  kordelfrance/Olfaction-Vision-Language-Embeddings

Each model combines:
  - A CLIP vision-language backbone (openai/clip-vit-base-patch32)
  - A TorchScript olfactory encoder
  - A TorchScript graph neural network (base) or graph-attention network (gat)

The result is a joint embedding that aligns olfaction, vision, and language
in a shared latent space.
"""
from __future__ import annotations

from typing import List, Union

import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download
from PIL import Image
from torchvision import transforms
from transformers import CLIPModel

_HF_REPO = "kordelfrance/Olfaction-Vision-Language-Embeddings"
_CLIP_BACKBONE = "openai/clip-vit-base-patch32"

# Maps public variant name → file paths inside the HF repo and config.
_VARIANTS: dict = {
    "colip-small-base": {
        "embed_dim": 512,
        "encoder_file": "model/colip-small/base/olf_encoder_512.pt",
        "gnn_file": "model/colip-small/base/gnn_512.pt",
        "is_gat": False,
    },
    "colip-small-gat": {
        "embed_dim": 512,
        "encoder_file": "model/colip-small/gat/olf_encoder_gat_512.pt",
        "gnn_file": "model/colip-small/gat/gnn_gat_512.pt",
        "is_gat": True,
    },
    "colip-large-base": {
        "embed_dim": 2048,
        "encoder_file": "model/colip-large/base/olf_encoder_2048.pt",
        "gnn_file": "model/colip-large/base/gnn_2048.pt",
        "is_gat": False,
    },
    "colip-large-gat": {
        "embed_dim": 2048,
        "encoder_file": "model/colip-large/gat/olf_encoder_gat_2048.pt",
        "gnn_file": "model/colip-large/gat/gnn_gat_2048.pt",
        "is_gat": True,
    },
}

_IMAGE_TRANSFORM = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])


class ColipModel:
    """
    COLIP: Olfaction-Vision-Language Embedding model.

    Produces joint olfaction-vision-language embeddings by fusing a CLIP
    vision backbone with a trained olfactory encoder and a graph neural
    network (base variants) or graph-attention network (gat variants).

    Four variants are available:

    =================== ========== ============ ================================
    Variant             Embed dim  Architecture Best for
    =================== ========== ============ ================================
    ``colip-small-base``   512     base GNN     Fast inference / edge devices
    ``colip-small-gat``    512     GAT          Higher accuracy on edge devices
    ``colip-large-base``  2048     base GNN     Accuracy-critical tasks
    ``colip-large-gat``   2048     GAT          Highest accuracy, slower
    =================== ========== ============ ================================

    Model weights are downloaded automatically from HuggingFace Hub and cached
    locally on first use.

    Example — single embedding::

        import scentience as scn
        from PIL import Image

        model = scn.ColipModel.from_pretrained("colip-small-base")

        image   = Image.open("scene.jpg")
        olf_vec = [0.12, 0.0, 0.85, ...]   # 138-dimensional descriptor

        embedding = model.embed(image, olf_vec)   # torch.Tensor (512,)

    Example — NumPy output::

        arr = model.embed_numpy(image, olf_vec)   # np.ndarray (512,)

    Example — large GAT variant::

        model = scn.ColipModel.from_pretrained("colip-large-gat")
        embedding = model.embed(image, olf_vec)   # torch.Tensor (2048,)
    """

    HF_REPO: str = _HF_REPO
    CLIP_BACKBONE: str = _CLIP_BACKBONE

    def __init__(
        self,
        variant: str = "colip-small-base",
        device: str | None = None,
    ) -> None:
        """
        Load a COLIP model. Downloads weights from HuggingFace Hub if not
        already cached locally.

        Args:
            variant: Model variant name. One of:
                     ``"colip-small-base"``, ``"colip-small-gat"``,
                     ``"colip-large-base"``, ``"colip-large-gat"``.
            device:  Torch device string (``"cpu"``, ``"cuda"``, ``"mps"``).
                     Auto-detects CUDA → MPS → CPU when omitted.

        Raises:
            ValueError: If *variant* is not recognised.
        """
        if variant not in _VARIANTS:
            raise ValueError(
                f"Unknown variant '{variant}'. "
                f"Choose from: {', '.join(_VARIANTS)}"
            )

        cfg = _VARIANTS[variant]
        self.variant: str = variant
        self.embed_dim: int = cfg["embed_dim"]
        self._is_gat: bool = cfg["is_gat"]
        self.device: torch.device = _resolve_device(device)

        # --- CLIP vision-language backbone ---
        self._clip: CLIPModel = (
            CLIPModel.from_pretrained(self.CLIP_BACKBONE).to(self.device)
        )
        self._clip.eval()

        # --- TorchScript olfactory encoder ---
        encoder_path = hf_hub_download(
            repo_id=self.HF_REPO, filename=cfg["encoder_file"]
        )
        self._olf_encoder = torch.jit.load(
            encoder_path, map_location=self.device
        )
        self._olf_encoder.eval()

        # --- TorchScript GNN / GAT ---
        gnn_path = hf_hub_download(
            repo_id=self.HF_REPO, filename=cfg["gnn_file"]
        )
        self._gnn = torch.jit.load(gnn_path, map_location=self.device)
        self._gnn.eval()

        # CLIP base-patch32 outputs 512-dim vectors. Large COLIP variants
        # operate in 2048-dim space, so we add a learned projection.
        clip_proj_dim: int = self._clip.config.projection_dim  # 512
        if self.embed_dim != clip_proj_dim:
            self._vision_proj: nn.Module | None = nn.Linear(
                clip_proj_dim, self.embed_dim
            ).to(self.device)
        else:
            self._vision_proj = None

    # ------------------------------------------------------------------
    # Construction helpers
    # ------------------------------------------------------------------

    @classmethod
    def from_pretrained(
        cls,
        variant: str = "colip-small-base",
        device: str | None = None,
    ) -> "ColipModel":
        """
        Load a COLIP model from HuggingFace Hub.

        Mirrors the standard HuggingFace ``from_pretrained`` convention.
        Equivalent to ``ColipModel(variant, device)``.

        Args:
            variant: Model variant. One of ``"colip-small-base"``,
                     ``"colip-small-gat"``, ``"colip-large-base"``,
                     ``"colip-large-gat"``.
            device:  Torch device string. Defaults to auto-detect.

        Returns:
            ColipModel ready for inference.
        """
        return cls(variant=variant, device=device)

    # ------------------------------------------------------------------
    # Inference
    # ------------------------------------------------------------------

    def embed(
        self,
        image: Image.Image,
        olf_vec: Union[List[float], "torch.Tensor"],
    ) -> torch.Tensor:
        """
        Compute a joint olfaction-vision-language embedding.

        Args:
            image:   PIL Image of the scene to embed.
            olf_vec: Olfactory descriptor vector — a 138-dimensional sequence
                     of floats (list, NumPy array, or 1-D torch.Tensor).

        Returns:
            torch.Tensor: 1-D embedding of shape ``(embed_dim,)``.
        """
        img_tensor = _IMAGE_TRANSFORM(image).unsqueeze(0).to(self.device)

        if not isinstance(olf_vec, torch.Tensor):
            olf_tensor = torch.tensor(olf_vec, dtype=torch.float32)
        else:
            olf_tensor = olf_vec.float()
        olf_tensor = olf_tensor.unsqueeze(0).to(self.device)

        with torch.no_grad():
            vision_embed = self._clip.get_image_features(pixel_values=img_tensor)

            if self._vision_proj is not None:
                vision_embed = self._vision_proj(vision_embed)

            olf_embed = self._olf_encoder(olf_tensor)

            if self._is_gat:
                # Build a fully-connected 2-node graph: [vision, olfaction]
                nodes = torch.cat([vision_embed, olf_embed], dim=0)
                n = nodes.size(0)
                edge_index = torch.cartesian_prod(
                    torch.arange(n, device=self.device),
                    torch.arange(n, device=self.device),
                ).T
                result = self._gnn(nodes, edge_index)
            else:
                result = self._gnn(vision_embed, olf_embed)

        return result.squeeze()

    def embed_numpy(
        self,
        image: Image.Image,
        olf_vec: Union[List[float], "torch.Tensor"],
    ) -> "object":  # np.ndarray
        """
        Like :meth:`embed` but returns a NumPy array on CPU.

        Args:
            image:   PIL Image.
            olf_vec: Olfactory descriptor vector (138-dimensional).

        Returns:
            np.ndarray: 1-D array of shape ``(embed_dim,)``.
        """
        return self.embed(image, olf_vec).cpu().numpy()

    # ------------------------------------------------------------------
    # Dunder
    # ------------------------------------------------------------------

    def __repr__(self) -> str:
        return (
            f"ColipModel(variant='{self.variant}', "
            f"embed_dim={self.embed_dim}, "
            f"device='{self.device}')"
        )


# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------

def _resolve_device(device: str | None) -> torch.device:
    if device is not None:
        return torch.device(device)
    if torch.cuda.is_available():
        return torch.device("cuda")
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")
