"""
Multimodal model adapter — load CLIP, LLaVA, BLIP, and other multimodal models.

Supports:
- CLIP (OpenAI) — contrastive text-image models
- LLaVA — large language and vision assistant
- BLIP/BLIP2 — bootstrapped language-image pretraining
- Other HuggingFace multimodal models

Usage::

    from model_osmosis.adapters.multimodal import MultimodalAdapter

    adapter = MultimodalAdapter()

    # Load CLIP
    clip = adapter.load("openai/clip-vit-base-patch32")

    # Encode text and images
    text_features = clip.encode_text(input_ids)
    image_features = clip.encode_image(pixel_values)

    # Load LLaVA
    llava = adapter.load("llava-hf/llava-1.5-7b-hf")
"""

from __future__ import annotations

from enum import Enum
from typing import Any, Dict, List, Optional, Union

import torch
import torch.nn as nn
import numpy as np


class MultimodalModelType(Enum):
    """Types of multimodal models."""

    CLIP = "clip"
    LLAVA = "llava"
    BLIP = "blip"
    BLIP2 = "blip2"
    GIT = "git"
    GENERIC = "generic"


class MultimodalExpert(nn.Module):
    """Wrapped multimodal model supporting multiple input types.

    Args:
        model: The multimodal model.
        hidden_size: Common output dimension.
        name: Human-readable name.
        text_dim: Hidden dimension for text features.
        vision_dim: Hidden dimension for vision features.
        model_type: Type of multimodal model.
    """

    def __init__(
        self,
        model: nn.Module,
        hidden_size: int,
        name: str = "multimodal_expert",
        text_dim: Optional[int] = None,
        vision_dim: Optional[int] = None,
        model_type: MultimodalModelType = MultimodalModelType.GENERIC,
    ) -> None:
        super().__init__()
        self.model = model
        self.hidden_size = hidden_size
        self.feature_dim = hidden_size
        self.name = name
        self.text_dim = text_dim or hidden_size
        self.vision_dim = vision_dim or hidden_size
        self.model_type = model_type

        # Freeze weights
        for param in self.model.parameters():
            param.requires_grad = False

    def encode_text(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Encode text to feature space.

        Args:
            input_ids: Token IDs (batch, seq_len).
            attention_mask: Optional attention mask.

        Returns:
            Text features (batch, text_dim).
        """
        with torch.no_grad():
            if self.model_type == MultimodalModelType.CLIP:
                kwargs = {"input_ids": input_ids}
                if attention_mask is not None:
                    kwargs["attention_mask"] = attention_mask
                return self.model.get_text_features(**kwargs)
            else:
                # Generic: use text_model if available
                if hasattr(self.model, "text_model"):
                    outputs = self.model.text_model(
                        input_ids=input_ids,
                        attention_mask=attention_mask,
                    )
                    if hasattr(outputs, "pooler_output") and outputs.pooler_output is not None:
                        return outputs.pooler_output
                    return outputs.last_hidden_state[:, 0, :]
                else:
                    outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
                    if hasattr(outputs, "text_embeds"):
                        return outputs.text_embeds
                    return outputs[0]

    def encode_image(self, pixel_values: torch.Tensor) -> torch.Tensor:
        """Encode images to feature space.

        Args:
            pixel_values: Image tensor (batch, channels, height, width).

        Returns:
            Image features (batch, vision_dim).
        """
        with torch.no_grad():
            if self.model_type == MultimodalModelType.CLIP:
                return self.model.get_image_features(pixel_values=pixel_values)
            else:
                # Generic: use vision_model if available
                if hasattr(self.model, "vision_model"):
                    outputs = self.model.vision_model(pixel_values=pixel_values)
                    if hasattr(outputs, "pooler_output") and outputs.pooler_output is not None:
                        return outputs.pooler_output
                    return outputs.last_hidden_state[:, 0, :]
                else:
                    outputs = self.model(pixel_values=pixel_values)
                    if hasattr(outputs, "image_embeds"):
                        return outputs.image_embeds
                    return outputs[0]

    def forward(
        self,
        x: torch.Tensor,
        modality: str = "vision",
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Forward pass for specified modality.

        Args:
            x: Input tensor (format depends on modality).
            modality: Which modality to use ('text' or 'vision').
            attention_mask: Optional attention mask for text.

        Returns:
            Feature tensor.
        """
        if modality == "text":
            return self.encode_text(x, attention_mask)
        else:
            return self.encode_image(x)

    def compute_similarity(
        self,
        text_features: torch.Tensor,
        image_features: torch.Tensor,
    ) -> torch.Tensor:
        """Compute similarity between text and image features.

        Args:
            text_features: Text embeddings (batch_t, dim).
            image_features: Image embeddings (batch_i, dim).

        Returns:
            Similarity matrix (batch_t, batch_i).
        """
        # Normalize
        text_features = text_features / text_features.norm(dim=-1, keepdim=True)
        image_features = image_features / image_features.norm(dim=-1, keepdim=True)

        # Cosine similarity
        return text_features @ image_features.T

    def extra_repr(self) -> str:
        return (
            f"name={self.name}, hidden_size={self.hidden_size}, "
            f"text_dim={self.text_dim}, vision_dim={self.vision_dim}, "
            f"model_type={self.model_type.value}"
        )


class MultimodalAdapter:
    """Load multimodal models from HuggingFace.

    Args:
        device: Device to load models to.
        torch_dtype: Data type for model weights.
        verbose: Print loading progress.
    """

    # Model type detection patterns
    MODEL_PATTERNS = {
        MultimodalModelType.CLIP: ["clip", "openclip"],
        MultimodalModelType.LLAVA: ["llava"],
        MultimodalModelType.BLIP2: ["blip2", "blip-2"],
        MultimodalModelType.BLIP: ["blip"],
        MultimodalModelType.GIT: ["git"],
    }

    def __init__(
        self,
        device: str = "cpu",
        torch_dtype: Optional[torch.dtype] = None,
        verbose: bool = True,
    ) -> None:
        self.device = device
        self.torch_dtype = torch_dtype
        self.verbose = verbose

    def load(
        self,
        model_name: str,
        name: Optional[str] = None,
        **kwargs,
    ) -> MultimodalExpert:
        """Load a multimodal model.

        Args:
            model_name: HuggingFace model ID.
            name: Human-readable name.

        Returns:
            MultimodalExpert ready for stacking.
        """
        try:
            from transformers import AutoModel, AutoConfig
        except ImportError:
            raise ImportError(
                "Multimodal adapter requires `transformers`. "
                "Install with: pip install transformers"
            )

        if self.verbose:
            print(f"[Multimodal] Loading '{model_name}'...")

        # Detect model type
        model_type = self._detect_type(model_name)

        # Load config
        config = AutoConfig.from_pretrained(model_name)

        # Get dimensions
        hidden_size, text_dim, vision_dim = self._get_dimensions(config, model_name)

        # Load model
        load_kwargs = dict(kwargs)
        if self.torch_dtype:
            load_kwargs["torch_dtype"] = self.torch_dtype

        if model_type == MultimodalModelType.CLIP:
            from transformers import CLIPModel
            model = CLIPModel.from_pretrained(model_name, **load_kwargs)
        elif model_type == MultimodalModelType.LLAVA:
            from transformers import LlavaForConditionalGeneration
            model = LlavaForConditionalGeneration.from_pretrained(model_name, **load_kwargs)
        elif model_type == MultimodalModelType.BLIP:
            from transformers import BlipModel
            model = BlipModel.from_pretrained(model_name, **load_kwargs)
        elif model_type == MultimodalModelType.BLIP2:
            from transformers import Blip2Model
            model = Blip2Model.from_pretrained(model_name, **load_kwargs)
        else:
            model = AutoModel.from_pretrained(model_name, **load_kwargs)

        model = model.to(self.device)
        model.eval()

        if self.verbose:
            total_params = sum(p.numel() for p in model.parameters())
            print(
                f"[Multimodal] Loaded '{name or model_name.split('/')[-1]}': "
                f"{total_params:,} params, hidden_size={hidden_size}"
            )

        return MultimodalExpert(
            model=model,
            hidden_size=hidden_size,
            name=name or model_name.split("/")[-1],
            text_dim=text_dim,
            vision_dim=vision_dim,
            model_type=model_type,
        )

    def _detect_type(self, model_name: str) -> MultimodalModelType:
        """Detect multimodal model type from name."""
        name_lower = model_name.lower()
        for model_type, patterns in self.MODEL_PATTERNS.items():
            for pattern in patterns:
                if pattern in name_lower:
                    return model_type
        return MultimodalModelType.GENERIC

    def _get_dimensions(
        self, config: Any, model_name: str
    ) -> tuple:
        """Get hidden dimensions from config."""
        text_dim = None
        vision_dim = None
        hidden_size = None

        # Check nested configs
        if hasattr(config, "text_config"):
            text_dim = getattr(config.text_config, "hidden_size", None)
        if hasattr(config, "vision_config"):
            vision_dim = getattr(config.vision_config, "hidden_size", None)

        # Check projection dim (CLIP-style)
        if hasattr(config, "projection_dim"):
            hidden_size = config.projection_dim
        elif text_dim and vision_dim:
            hidden_size = max(text_dim, vision_dim)
        else:
            hidden_size = getattr(config, "hidden_size", 768)

        text_dim = text_dim or hidden_size
        vision_dim = vision_dim or hidden_size

        return hidden_size, text_dim, vision_dim

    def load_multiple(
        self,
        model_names: List[str],
        names: Optional[List[str]] = None,
    ) -> List[MultimodalExpert]:
        """Load multiple multimodal models.

        Args:
            model_names: List of HuggingFace model IDs.
            names: Optional list of names.

        Returns:
            List of MultimodalExpert modules.
        """
        names = names or [None] * len(model_names)
        experts = []
        for model_name, name in zip(model_names, names):
            expert = self.load(model_name, name=name)
            experts.append(expert)
        return experts
