"""
Ultralytics YOLO adapter.

Supports YOLOv8, YOLOv10, YOLOv11 and other Ultralytics-based models.
"""

from typing import List, Dict, Any, Tuple, Optional
from pathlib import Path
import torch
import torch.nn as nn
import numpy as np
from PIL import Image

from ..core.base import BaseModelAdapter, Detection, FPNLayers


class YOLOAdapter(BaseModelAdapter):
    """
    Adapter for Ultralytics YOLO models.
    
    Tested with: YOLOv8, YOLOv10, YOLOv11
    """
    
    # Layer names for different YOLO versions (may need adjustment)
    # These target the neck/FPN outputs
    _FPN_LAYER_PATTERNS = {
        'p3': ['model.model.15', 'model.15'],  # Small objects, high res
        'p4': ['model.model.18', 'model.18'],  # Medium objects
        'p5': ['model.model.21', 'model.21'],  # Large objects, low res
    }
    
    def __init__(self, model_path: Optional[str] = None):
        """
        Initialize YOLO adapter.
        
        Args:
            model_path: Path to .pt weights or model name (e.g., 'yolov8n.pt')
        """
        self._model = None
        self._device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self._fpn_layers: Dict[str, nn.Module] = {}
        
        if model_path:
            self.load(model_path)
    
    def load(self, path: str) -> None:
        """Load YOLO model from path."""
        from ultralytics import YOLO
        self._model = YOLO(path)
        self._model.to(self._device)
        self._discover_fpn_layers()
    
    def _discover_fpn_layers(self) -> None:
        """Find FPN layers in the model architecture."""
        pytorch_model = self.get_pytorch_model()
        
        for scale, patterns in self._FPN_LAYER_PATTERNS.items():
            for pattern in patterns:
                try:
                    layer = self._get_module_by_name(pytorch_model, pattern)
                    if layer is not None:
                        self._fpn_layers[scale] = layer
                        break
                except (AttributeError, KeyError):
                    continue
        
        # Fallback: find last few C2f/Conv layers if patterns don't match
        if len(self._fpn_layers) < 3:
            self._fallback_fpn_discovery()
    
    def _get_module_by_name(self, model: nn.Module, name: str) -> nn.Module:
        """Get nested module by dot-separated name."""
        parts = name.split('.')
        module = model
        for part in parts:
            if part.isdigit():
                module = module[int(part)]
            else:
                module = getattr(module, part)
        return module
    
    def _fallback_fpn_discovery(self) -> None:
        """Fallback: find suitable layers by type inspection."""
        pytorch_model = self.get_pytorch_model()
        c2f_layers = []
        
        for name, module in pytorch_model.named_modules():
            # C2f is the main feature block in YOLOv8+
            if 'C2f' in type(module).__name__ or 'SPPF' in type(module).__name__:
                c2f_layers.append((name, module))
        
        # Take last 3 suitable layers as P3, P4, P5
        if len(c2f_layers) >= 3:
            self._fpn_layers = {
                'p5': c2f_layers[-1][1],
                'p4': c2f_layers[-2][1],
                'p3': c2f_layers[-3][1],
            }
    
    def get_pytorch_model(self) -> nn.Module:
        """Return underlying PyTorch model."""
        if self._model is None:
            raise RuntimeError("Model not loaded. Call load() first.")
        return self._model.model
    
    def get_fpn_layers(self) -> Dict[str, nn.Module]:
        """Return FPN layer modules for hooking."""
        if not self._fpn_layers:
            raise RuntimeError("FPN layers not discovered. Is model loaded?")
        return self._fpn_layers
    
    def get_detection_head(self) -> nn.Module:
        """Return detection head module."""
        pytorch_model = self.get_pytorch_model()
        # Detection head is typically the last module
        if hasattr(pytorch_model, 'model'):
            return pytorch_model.model[-1]
        return list(pytorch_model.children())[-1]
    
    def predict(self, image: Any, conf: float = 0.25) -> List[Detection]:
        """
        Run inference and return detections.
        
        Args:
            image: Path, numpy array, or PIL Image
            conf: Confidence threshold
            
        Returns:
            List of Detection objects
        """
        if self._model is None:
            raise RuntimeError("Model not loaded. Call load() first.")
        
        results = self._model(image, conf=conf, verbose=False)
        
        detections = []
        for result in results:
            if result.boxes is None:
                continue
            
            boxes = result.boxes
            for i in range(len(boxes)):
                bbox = boxes.xyxy[i].cpu().numpy()
                conf_score = float(boxes.conf[i].cpu())
                class_id = int(boxes.cls[i].cpu())
                class_name = self._model.names[class_id]
                
                detections.append(Detection(
                    bbox=tuple(bbox),
                    confidence=conf_score,
                    class_id=class_id,
                    class_name=class_name
                ))
        
        return detections
    
    def get_class_names(self) -> List[str]:
        """Return list of class names."""
        if self._model is None:
            raise RuntimeError("Model not loaded. Call load() first.")
        return list(self._model.names.values())
    
    def preprocess(self, image: Any) -> torch.Tensor:
        """
        Preprocess image for model input.
        
        Args:
            image: Path, numpy array, or PIL Image
            
        Returns:
            Tensor of shape (1, 3, H, W) normalized to [0, 1]
        """
        if isinstance(image, (str, Path)):
            image = Image.open(image).convert('RGB')
        elif isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        
        # Get model input size
        input_size = self.get_input_size()
        image = image.resize((input_size[1], input_size[0]))
        
        # Convert to tensor
        tensor = torch.from_numpy(np.array(image)).float()
        tensor = tensor.permute(2, 0, 1)  # HWC -> CHW
        tensor = tensor.unsqueeze(0)  # Add batch dim
        tensor = tensor / 255.0  # Normalize to [0, 1]
        
        return tensor.to(self._device)
    
    def get_input_size(self) -> Tuple[int, int]:
        """Return expected input size (height, width)."""
        # Default YOLO input size, may vary by model
        return (640, 640)
    
    def forward_with_hooks(
        self, 
        image: Any,
        target_layers: Optional[List[str]] = None
    ) -> Tuple[List[Detection], Dict[str, torch.Tensor]]:
        """
        Forward pass capturing intermediate features.
        
        Args:
            image: Input image
            target_layers: Layer names to capture. None = capture FPN layers.
            
        Returns:
            (detections, features_dict)
        """
        from ..core.hooks import ActivationHook
        
        if target_layers is None:
            layers_to_hook = self.get_fpn_layers()
        else:
            pytorch_model = self.get_pytorch_model()
            layers_to_hook = {
                name: self._get_module_by_name(pytorch_model, name)
                for name in target_layers
            }
        
        hooks = {name: ActivationHook(module) for name, module in layers_to_hook.items()}
        
        try:
            detections = self.predict(image)
            features = {name: hook.activation for name, hook in hooks.items()}
            return detections, features
        finally:
            for hook in hooks.values():
                hook.remove()
    
    @property
    def device(self) -> torch.device:
        """Return device model is on."""
        return self._device
