"""
GridCellTrace - Visualize grid cell activations for anchor-free YOLO.

For YOLO v8+ which uses anchor-free detection with FPN,
this visualizes which grid cells activated across P3/P4/P5 scales.
"""

from typing import Any, Dict, List, Tuple, Optional
from dataclasses import dataclass
import torch
import torch.nn.functional as F
import numpy as np

from ..core.base import BaseModelAdapter, Detection
from ..core.hooks import ActivationHook


@dataclass
class GridCell:
    """Single grid cell that contributed to detection."""
    scale: str  # 'p3', 'p4', or 'p5'
    x: int  # Grid x coordinate
    y: int  # Grid y coordinate
    confidence: float
    class_id: Optional[int] = None


@dataclass
class GridTraceResult:
    """Result from GridCellTrace analysis."""
    
    # Per-scale activation maps (raw confidence at each cell)
    p3_activations: torch.Tensor  # (H3, W3) e.g., 80x80
    p4_activations: torch.Tensor  # (H4, W4) e.g., 40x40
    p5_activations: torch.Tensor  # (H5, W5) e.g., 20x20
    
    # Which cells produced final detections (after NMS)
    winning_cells: List[GridCell]
    
    # High-confidence cells that were suppressed by NMS
    runner_up_cells: List[GridCell]
    
    # Fused visualization at input resolution
    fused_activation: torch.Tensor
    
    # Scale statistics
    detections_per_scale: Dict[str, int]


class GridCellTrace:
    """
    Visualize grid cell activations for anchor-free YOLO.
    
    YOLO v8+ doesn't use traditional anchors. Instead, each cell
    in the FPN output predicts object presence directly.
    
    This traces:
    - Which cells activate at each FPN scale
    - Which cells "won" to produce final detections
    - Which cells were close but suppressed by NMS
    
    Usage:
        tracer = GridCellTrace(adapter)
        result = tracer.trace(image)
    """
    
    def __init__(self, adapter: BaseModelAdapter):
        self.adapter = adapter
        self.nms_threshold = 0.45  # Typical YOLO NMS threshold
    
    def trace(self, image: Any) -> GridTraceResult:
        """
        Trace grid cell activations for an image.
        
        Args:
            image: Input image
            
        Returns:
            GridTraceResult with activation maps and cell analysis
        """
        # Get detections and FPN features
        detections, fpn_features = self.adapter.forward_with_hooks(image)
        
        # Process each scale
        scale_activations = {}
        all_cells = []
        
        for scale in ['p3', 'p4', 'p5']:
            if scale in fpn_features and fpn_features[scale] is not None:
                features = fpn_features[scale]
                
                # Extract objectness/confidence from features
                activation_map = self._extract_objectness(features)
                scale_activations[scale] = activation_map
                
                # Find active cells
                cells = self._find_active_cells(activation_map, scale)
                all_cells.extend(cells)
        
        # Match cells to detections
        winning_cells = self._match_cells_to_detections(all_cells, detections, image)
        
        # Find runner-ups (high confidence but not in final detections)
        runner_ups = self._find_runner_ups(all_cells, winning_cells)
        
        # Fuse activations to input resolution
        input_size = self.adapter.get_input_size()
        fused = self._fuse_activations(scale_activations, input_size)
        
        # Count detections per scale
        dets_per_scale = self._count_detections_per_scale(winning_cells)
        
        return GridTraceResult(
            p3_activations=scale_activations.get('p3', torch.zeros(1)),
            p4_activations=scale_activations.get('p4', torch.zeros(1)),
            p5_activations=scale_activations.get('p5', torch.zeros(1)),
            winning_cells=winning_cells,
            runner_up_cells=runner_ups,
            fused_activation=fused,
            detections_per_scale=dets_per_scale,
        )
    
    def _extract_objectness(self, features: torch.Tensor) -> torch.Tensor:
        """
        Extract objectness/confidence map from feature tensor.
        
        For YOLO, this typically comes from detection head output
        or can be approximated from feature magnitude.
        """
        if features is None:
            return torch.zeros(1)
        
        # features shape: (B, C, H, W)
        # Use L2 norm across channels as proxy for "activation strength"
        if len(features.shape) == 4:
            activation = (features ** 2).sum(dim=1).sqrt()  # (B, H, W)
            activation = activation.squeeze(0)  # (H, W)
        else:
            activation = features.squeeze()
        
        # Normalize to [0, 1]
        act_min = activation.min()
        act_max = activation.max()
        if act_max - act_min > 1e-8:
            activation = (activation - act_min) / (act_max - act_min)
        
        return activation.detach().cpu()
    
    def _find_active_cells(
        self, 
        activation_map: torch.Tensor, 
        scale: str,
        threshold: float = 0.3
    ) -> List[GridCell]:
        """Find grid cells with activation above threshold."""
        cells = []
        
        if activation_map.numel() <= 1:
            return cells
        
        for y in range(activation_map.shape[0]):
            for x in range(activation_map.shape[1]):
                conf = float(activation_map[y, x])
                if conf > threshold:
                    cells.append(GridCell(
                        scale=scale,
                        x=x,
                        y=y,
                        confidence=conf
                    ))
        
        # Sort by confidence
        cells.sort(key=lambda c: c.confidence, reverse=True)
        
        return cells
    
    def _match_cells_to_detections(
        self,
        cells: List[GridCell],
        detections: List[Detection],
        image: Any
    ) -> List[GridCell]:
        """
        Match grid cells to final detections based on location.
        """
        winning_cells = []
        input_size = self.adapter.get_input_size()
        
        for det in detections:
            # Get detection center
            cx = (det.bbox[0] + det.bbox[2]) / 2
            cy = (det.bbox[1] + det.bbox[3]) / 2
            
            # Determine likely scale based on object size
            obj_size = max(det.bbox[2] - det.bbox[0], det.bbox[3] - det.bbox[1])
            relative_size = obj_size / max(input_size)
            
            if relative_size < 0.1:
                likely_scale = 'p3'  # Small objects
                grid_size = 80
            elif relative_size < 0.3:
                likely_scale = 'p4'  # Medium objects
                grid_size = 40
            else:
                likely_scale = 'p5'  # Large objects
                grid_size = 20
            
            # Find matching cell
            grid_x = int(cx / input_size[1] * grid_size)
            grid_y = int(cy / input_size[0] * grid_size)
            
            # Find cell that matches
            for cell in cells:
                if (cell.scale == likely_scale and 
                    abs(cell.x - grid_x) <= 1 and 
                    abs(cell.y - grid_y) <= 1):
                    cell.class_id = det.class_id
                    winning_cells.append(cell)
                    break
            else:
                # Create synthetic cell if not found
                winning_cells.append(GridCell(
                    scale=likely_scale,
                    x=grid_x,
                    y=grid_y,
                    confidence=det.confidence,
                    class_id=det.class_id
                ))
        
        return winning_cells
    
    def _find_runner_ups(
        self,
        all_cells: List[GridCell],
        winning_cells: List[GridCell],
        max_count: int = 20
    ) -> List[GridCell]:
        """Find high-confidence cells that didn't produce final detections."""
        winning_set = {(c.scale, c.x, c.y) for c in winning_cells}
        
        runner_ups = [
            c for c in all_cells 
            if (c.scale, c.x, c.y) not in winning_set
            and c.confidence > 0.5
        ]
        
        # Sort by confidence and take top N
        runner_ups.sort(key=lambda c: c.confidence, reverse=True)
        return runner_ups[:max_count]
    
    def _fuse_activations(
        self,
        scale_activations: Dict[str, torch.Tensor],
        target_size: Tuple[int, int]
    ) -> torch.Tensor:
        """Fuse multi-scale activations to single map."""
        fused = torch.zeros(target_size)
        
        weights = {'p3': 1.0, 'p4': 1.0, 'p5': 1.0}  # Equal weights
        
        for scale, activation in scale_activations.items():
            if activation.numel() > 1:
                # Upsample to target size
                upsampled = F.interpolate(
                    activation.unsqueeze(0).unsqueeze(0),
                    size=target_size,
                    mode='bilinear',
                    align_corners=False
                ).squeeze()
                
                fused = fused + weights[scale] * upsampled
        
        # Normalize
        fused_min = fused.min()
        fused_max = fused.max()
        if fused_max - fused_min > 1e-8:
            fused = (fused - fused_min) / (fused_max - fused_min)
        
        return fused
    
    def _count_detections_per_scale(
        self, 
        winning_cells: List[GridCell]
    ) -> Dict[str, int]:
        """Count how many detections came from each scale."""
        counts = {'p3': 0, 'p4': 0, 'p5': 0}
        for cell in winning_cells:
            if cell.scale in counts:
                counts[cell.scale] += 1
        return counts
