"""
Feature extraction utilities.

Helper functions for extracting and processing intermediate features.
"""

from typing import Dict, List, Optional
import torch
import torch.nn as nn
import torch.nn.functional as F


class FeatureExtractor:
    """
    Extract features from multiple layers simultaneously.
    
    Usage:
        extractor = FeatureExtractor(model, ['layer1', 'layer2.conv'])
        features = extractor.extract(input_tensor)
        # features = {'layer1': Tensor, 'layer2.conv': Tensor}
    """
    
    def __init__(self, model: nn.Module, layer_names: List[str]):
        self.model = model
        self.layer_names = layer_names
        self._features: Dict[str, torch.Tensor] = {}
        self._hooks = []
        self._register_hooks()
    
    def _get_layer(self, name: str) -> nn.Module:
        """Get layer by dot-separated name."""
        module = self.model
        for part in name.split('.'):
            if part.isdigit():
                module = module[int(part)]
            else:
                module = getattr(module, part)
        return module
    
    def _register_hooks(self):
        """Register forward hooks on target layers."""
        for name in self.layer_names:
            layer = self._get_layer(name)
            handle = layer.register_forward_hook(
                lambda mod, inp, out, n=name: self._save_feature(n, out)
            )
            self._hooks.append(handle)
    
    def _save_feature(self, name: str, output: torch.Tensor):
        """Store feature map."""
        self._features[name] = output.detach()
    
    def extract(self, input_tensor: torch.Tensor) -> Dict[str, torch.Tensor]:
        """
        Run forward pass and return extracted features.
        
        Args:
            input_tensor: Model input
            
        Returns:
            Dict mapping layer names to their output tensors
        """
        self._features.clear()
        with torch.no_grad():
            self.model(input_tensor)
        return dict(self._features)
    
    def remove(self):
        """Remove all hooks."""
        for handle in self._hooks:
            handle.remove()
        self._hooks.clear()
    
    def __enter__(self):
        return self
    
    def __exit__(self, *args):
        self.remove()


def interpolate_to_size(
    feature_map: torch.Tensor, 
    target_size: tuple,
    mode: str = 'bilinear'
) -> torch.Tensor:
    """
    Interpolate feature map to target spatial size.
    
    Args:
        feature_map: Tensor of shape (B, C, H, W)
        target_size: (target_H, target_W)
        mode: Interpolation mode
        
    Returns:
        Interpolated tensor of shape (B, C, target_H, target_W)
    """
    return F.interpolate(
        feature_map, 
        size=target_size, 
        mode=mode, 
        align_corners=False if mode != 'nearest' else None
    )


def normalize_attribution(attribution: torch.Tensor) -> torch.Tensor:
    """
    Normalize attribution map to [0, 1] range.
    
    Args:
        attribution: Tensor of any shape
        
    Returns:
        Normalized tensor with same shape, values in [0, 1]
    """
    attr_min = attribution.min()
    attr_max = attribution.max()
    
    if attr_max - attr_min > 1e-8:
        return (attribution - attr_min) / (attr_max - attr_min)
    else:
        return torch.zeros_like(attribution)


def aggregate_channels(feature_map: torch.Tensor, method: str = 'mean') -> torch.Tensor:
    """
    Aggregate feature map across channel dimension.
    
    Args:
        feature_map: Tensor of shape (B, C, H, W)
        method: 'mean', 'max', or 'l2'
        
    Returns:
        Tensor of shape (B, 1, H, W)
    """
    if method == 'mean':
        return feature_map.mean(dim=1, keepdim=True)
    elif method == 'max':
        return feature_map.max(dim=1, keepdim=True)[0]
    elif method == 'l2':
        return (feature_map ** 2).sum(dim=1, keepdim=True).sqrt()
    else:
        raise ValueError(f"Unknown aggregation method: {method}")
