"""
Base classes for model adapters.

To add support for a new model architecture, extend BaseModelAdapter.
"""

from abc import ABC, abstractmethod
from typing import List, Any, Dict, Tuple
from dataclasses import dataclass
import torch


@dataclass
class Detection:
    """Single detection result."""
    bbox: Tuple[float, float, float, float]  # x1, y1, x2, y2
    confidence: float
    class_id: int
    class_name: str


@dataclass 
class FPNLayers:
    """Feature Pyramid Network layer outputs."""
    p3: torch.Tensor  # Small objects, high resolution (e.g., 80x80)
    p4: torch.Tensor  # Medium objects (e.g., 40x40)
    p5: torch.Tensor  # Large objects, low resolution (e.g., 20x20)


class BaseModelAdapter(ABC):
    """
    Abstract adapter for different model architectures.
    
    Extend this class to add support for new models (DETR, EfficientDet, etc.)
    """
    
    @abstractmethod
    def load(self, path: str) -> None:
        """Load model from file path or name."""
        pass
    
    @abstractmethod
    def get_pytorch_model(self) -> torch.nn.Module:
        """Return the underlying PyTorch model."""
        pass
    
    @abstractmethod
    def get_fpn_layers(self) -> Dict[str, torch.nn.Module]:
        """
        Return FPN layer modules for hooking.
        
        Returns:
            Dict with keys 'p3', 'p4', 'p5' mapping to layer modules.
        """
        pass
    
    @abstractmethod
    def get_detection_head(self) -> torch.nn.Module:
        """Return the detection head module."""
        pass
    
    @abstractmethod
    def predict(self, image: Any) -> List[Detection]:
        """
        Run inference and return detections.
        
        Args:
            image: Image path, numpy array, or PIL Image
            
        Returns:
            List of Detection objects
        """
        pass
    
    @abstractmethod
    def get_class_names(self) -> List[str]:
        """Return list of class names the model was trained on."""
        pass
    
    @abstractmethod
    def preprocess(self, image: Any) -> torch.Tensor:
        """
        Preprocess image for model input.
        
        Returns:
            Tensor of shape (1, 3, H, W) ready for model.
        """
        pass
    
    @abstractmethod
    def get_input_size(self) -> Tuple[int, int]:
        """Return expected input size (height, width)."""
        pass
    
    def forward_with_features(self, image: Any) -> Tuple[List[Detection], FPNLayers]:
        """
        Run forward pass and capture FPN features.
        
        Default implementation uses hooks. Override if model provides
        direct access to intermediate features.
        """
        from .hooks import ActivationHook
        
        fpn_modules = self.get_fpn_layers()
        hooks = {name: ActivationHook(module) for name, module in fpn_modules.items()}
        
        try:
            detections = self.predict(image)
            
            fpn_features = FPNLayers(
                p3=hooks['p3'].activation,
                p4=hooks['p4'].activation,
                p5=hooks['p5'].activation,
            )
            
            return detections, fpn_features
        finally:
            for hook in hooks.values():
                hook.remove()
