"""
Vision model adapter — load from HuggingFace, timm, or torchvision.

Supports:
- HuggingFace vision models (ViT, DINOv2, ConvNeXt, etc.)
- timm models (1000+ pretrained vision models)
- torchvision models (ResNet, EfficientNet, etc.)

Usage::

    from model_osmosis.adapters.vision import VisionAdapter

    adapter = VisionAdapter()

    # From HuggingFace (auto-detected)
    vit = adapter.load("google/vit-base-patch16-224", source="huggingface")

    # From timm
    resnet = adapter.load("resnet50", source="timm")

    # From torchvision
    efficientnet = adapter.load("efficientnet_b0", source="torchvision")

    # Auto-detect source
    model = adapter.load("google/vit-base-patch16-224")  # Detects HF
    model = adapter.load("resnet50")  # Detects timm
"""

from __future__ import annotations

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

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


class VisionSource(Enum):
    """Source for loading vision models."""

    HUGGINGFACE = "huggingface"
    TIMM = "timm"
    TORCHVISION = "torchvision"
    AUTO = "auto"


class VisionExpert(nn.Module):
    """Wrapped vision model as stackable expert.

    Args:
        model: The vision model backbone.
        feature_dim: Output feature dimension.
        name: Human-readable name.
        source: Where the model was loaded from.
        image_size: Expected input image size.
    """

    def __init__(
        self,
        model: nn.Module,
        feature_dim: int,
        name: str = "vision_expert",
        source: str = "unknown",
        image_size: int = 224,
    ) -> None:
        super().__init__()
        self.model = model
        self.feature_dim = feature_dim
        self.hidden_size = feature_dim  # Alias for compatibility
        self.name = name
        self.source = source
        self.image_size = image_size

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Extract features from images.

        Args:
            x: Image tensor (batch, 3, H, W) or embeddings (batch, seq, dim).

        Returns:
            Feature tensor (batch, feature_dim).
        """
        with torch.no_grad():
            features = self.model(x)

        # Handle different output types
        if isinstance(features, dict):
            features = features.get("pooler_output", features.get("last_hidden_state"))

        if features is None:
            raise ValueError("Model output is None")

        # Pool if needed
        if features.dim() == 3:
            features = features[:, 0, :]  # CLS token
        if features.dim() == 4:
            features = features.mean(dim=[2, 3])  # Global average pool

        return features

    def extra_repr(self) -> str:
        return (
            f"name={self.name}, feature_dim={self.feature_dim}, "
            f"source={self.source}, image_size={self.image_size}"
        )


class VisionAdapter:
    """Load vision models from multiple sources.

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

    # Common feature dimensions for quick lookup
    DIM_LOOKUP = {
        "vit-base": 768,
        "vit-large": 1024,
        "vit-huge": 1280,
        "dinov2-small": 384,
        "dinov2-base": 768,
        "dinov2-large": 1024,
        "dinov2-giant": 1536,
        "resnet18": 512,
        "resnet34": 512,
        "resnet50": 2048,
        "resnet101": 2048,
        "resnet152": 2048,
        "efficientnet_b0": 1280,
        "efficientnet_b1": 1280,
        "efficientnet_b2": 1408,
        "efficientnet_b3": 1536,
        "efficientnet_b4": 1792,
        "efficientnet_b5": 2048,
        "efficientnet_b6": 2304,
        "efficientnet_b7": 2560,
        "convnext_tiny": 768,
        "convnext_small": 768,
        "convnext_base": 1024,
        "convnext_large": 1536,
    }

    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,
        source: Union[str, VisionSource] = VisionSource.AUTO,
        name: Optional[str] = None,
        pretrained: bool = True,
        **kwargs,
    ) -> VisionExpert:
        """Load a vision model from any source.

        Args:
            model_name: Model identifier.
            source: Where to load from ('huggingface', 'timm', 'torchvision', 'auto').
            name: Human-readable name.
            pretrained: Load pretrained weights.

        Returns:
            VisionExpert ready for stacking.
        """
        if isinstance(source, str):
            source = VisionSource(source.lower())

        if source == VisionSource.AUTO:
            source = self._detect_source(model_name)

        if self.verbose:
            print(f"[Vision] Loading '{model_name}' from {source.value}...")

        if source == VisionSource.HUGGINGFACE:
            return self._load_from_huggingface(model_name, name, **kwargs)
        elif source == VisionSource.TIMM:
            return self._load_from_timm(model_name, name, pretrained, **kwargs)
        elif source == VisionSource.TORCHVISION:
            return self._load_from_torchvision(model_name, name, pretrained, **kwargs)

        raise ValueError(f"Unknown source: {source}")

    def _detect_source(self, model_name: str) -> VisionSource:
        """Auto-detect the best source for a model name."""
        # HuggingFace models have org/model format
        if "/" in model_name:
            return VisionSource.HUGGINGFACE

        # Check timm registry
        try:
            import timm
            if model_name in timm.list_models():
                return VisionSource.TIMM
        except ImportError:
            pass

        # Check torchvision
        try:
            import torchvision.models as tv_models
            if hasattr(tv_models, model_name):
                return VisionSource.TORCHVISION
        except ImportError:
            pass

        # Default to timm (largest model zoo)
        return VisionSource.TIMM

    def _load_from_huggingface(
        self,
        model_name: str,
        name: Optional[str],
        **kwargs,
    ) -> VisionExpert:
        """Load from HuggingFace Hub."""
        try:
            from transformers import AutoModel, AutoConfig
        except ImportError:
            raise ImportError(
                "HuggingFace loading requires `transformers`. "
                "Install with: pip install transformers"
            )

        config = AutoConfig.from_pretrained(model_name)
        model = AutoModel.from_pretrained(model_name, **kwargs)

        # Get feature dimension
        feature_dim = getattr(config, "hidden_size", None)
        if feature_dim is None:
            feature_dim = getattr(config, "embed_dim", 768)

        # Get image size
        image_size = getattr(config, "image_size", 224)
        if isinstance(image_size, (list, tuple)):
            image_size = image_size[0]

        model = model.to(self.device)
        if self.torch_dtype:
            model = model.to(self.torch_dtype)
        model.eval()

        if self.verbose:
            total_params = sum(p.numel() for p in model.parameters())
            print(f"[Vision] Loaded: {total_params:,} params, feature_dim={feature_dim}")

        return VisionExpert(
            model=model,
            feature_dim=feature_dim,
            name=name or model_name.split("/")[-1],
            source="huggingface",
            image_size=image_size,
        )

    def _load_from_timm(
        self,
        model_name: str,
        name: Optional[str],
        pretrained: bool,
        **kwargs,
    ) -> VisionExpert:
        """Load from timm."""
        try:
            import timm
        except ImportError:
            raise ImportError(
                "timm loading requires `timm`. "
                "Install with: pip install timm"
            )

        # Load model without classification head
        model = timm.create_model(
            model_name,
            pretrained=pretrained,
            num_classes=0,  # Remove classification head
            **kwargs,
        )

        feature_dim = model.num_features

        # Get image size from config
        cfg = model.pretrained_cfg if hasattr(model, "pretrained_cfg") else {}
        input_size = cfg.get("input_size", (3, 224, 224))
        image_size = input_size[-1] if isinstance(input_size, (list, tuple)) else 224

        model = model.to(self.device)
        if self.torch_dtype:
            model = model.to(self.torch_dtype)
        model.eval()

        if self.verbose:
            total_params = sum(p.numel() for p in model.parameters())
            print(f"[Vision] Loaded: {total_params:,} params, feature_dim={feature_dim}")

        return VisionExpert(
            model=model,
            feature_dim=feature_dim,
            name=name or model_name,
            source="timm",
            image_size=image_size,
        )

    def _load_from_torchvision(
        self,
        model_name: str,
        name: Optional[str],
        pretrained: bool,
        **kwargs,
    ) -> VisionExpert:
        """Load from torchvision."""
        try:
            import torchvision.models as tv_models
        except ImportError:
            raise ImportError(
                "torchvision loading requires `torchvision`. "
                "Install with: pip install torchvision"
            )

        # Get model class
        if not hasattr(tv_models, model_name):
            raise ValueError(f"Model '{model_name}' not found in torchvision.models")

        weights = "DEFAULT" if pretrained else None
        model = getattr(tv_models, model_name)(weights=weights, **kwargs)

        # Remove classification head and get feature dim
        feature_dim = self._remove_classifier(model)

        model = model.to(self.device)
        if self.torch_dtype:
            model = model.to(self.torch_dtype)
        model.eval()

        if self.verbose:
            total_params = sum(p.numel() for p in model.parameters())
            print(f"[Vision] Loaded: {total_params:,} params, feature_dim={feature_dim}")

        return VisionExpert(
            model=model,
            feature_dim=feature_dim,
            name=name or model_name,
            source="torchvision",
            image_size=224,
        )

    def _remove_classifier(self, model: nn.Module) -> int:
        """Remove classification head and return feature dimension."""
        # Try common head locations
        if hasattr(model, "fc"):
            feature_dim = model.fc.in_features
            model.fc = nn.Identity()
        elif hasattr(model, "classifier"):
            if isinstance(model.classifier, nn.Sequential):
                feature_dim = model.classifier[0].in_features
            elif isinstance(model.classifier, nn.Linear):
                feature_dim = model.classifier.in_features
            else:
                feature_dim = 512
            model.classifier = nn.Identity()
        elif hasattr(model, "head"):
            if isinstance(model.head, nn.Linear):
                feature_dim = model.head.in_features
            else:
                feature_dim = 512
            model.head = nn.Identity()
        elif hasattr(model, "heads"):
            if hasattr(model.heads, "head"):
                feature_dim = model.heads.head.in_features
            else:
                feature_dim = 512
            model.heads = nn.Identity()
        else:
            # Try to infer from last module
            feature_dim = 512
            for name, module in model.named_modules():
                if isinstance(module, nn.Linear):
                    feature_dim = module.in_features

        return feature_dim

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

        Args:
            model_names: List of model identifiers.
            sources: Optional list of sources per model.
            names: Optional list of names per model.

        Returns:
            List of VisionExpert modules.
        """
        sources = sources or ["auto"] * len(model_names)
        names = names or [None] * len(model_names)

        experts = []
        for model_name, source, name in zip(model_names, sources, names):
            expert = self.load(model_name, source=source, name=name)
            experts.append(expert)

        return experts

    def get_feature_dim(self, model_name: str, source: str = "auto") -> int:
        """Get feature dimension without loading full model.

        Args:
            model_name: Model identifier.
            source: Model source.

        Returns:
            Feature dimension.
        """
        # Check lookup table first
        name_lower = model_name.lower()
        for key, dim in self.DIM_LOOKUP.items():
            if key in name_lower:
                return dim

        # Try HuggingFace config
        if "/" in model_name:
            try:
                from transformers import AutoConfig
                config = AutoConfig.from_pretrained(model_name)
                return getattr(config, "hidden_size", 768)
            except Exception:
                pass

        # Default
        return 768

    def list_timm_models(self, filter: str = "") -> List[str]:
        """List available timm models.

        Args:
            filter: Optional filter string.

        Returns:
            List of model names.
        """
        try:
            import timm
            return timm.list_models(filter)
        except ImportError:
            return []
