"""
Production-grade reusable NMS module with YOLO-matching implementation.

This module provides class-specific and class-agnostic NMS implementations
that match YOLO's built-in behavior while being completely framework-agnostic.

Usage:
    from nms_module import AgnosticNMS
    
    nms = AgnosticNMS(iou_threshold=0.45, min_box_size=2.0)
    filtered_detections = nms.apply(detections, class_agnostic=True)
"""

from typing import Any, Dict, List, Optional, Tuple
import numpy as np

try:
    import torch
    from torchvision.ops import nms as torchvision_nms
    TORCHVISION_AVAILABLE = True
except ImportError:
    TORCHVISION_AVAILABLE = False


class AgnosticNMS:
    """
    Production-grade NMS implementation with YOLO-matching behavior.
    
    Features:
    - Class-specific and class-agnostic modes
    - Vectorized (PyTorch) and iterative fallback
    - Numerical stability enhancements
    - Box validation and filtering
    - Schema preservation
    - Zero side effects
    
    Attributes:
        iou_threshold: IoU threshold for suppression (default: 0.45)
        min_box_size: Minimum box width/height in pixels (default: 2.0)
        use_vectorized: Use torchvision.ops.nms if available (default: True)
        eps: Epsilon for numerical stability (default: 1e-7)
    """
    
    def __init__(
        self,
        iou_threshold: float = 0.45,
        min_box_size: float = 2.0,
        use_vectorized: bool = True,
        eps: float = 1e-7
    ):
        """
        Initialize NMS module.
        
        Args:
            iou_threshold: IoU threshold for suppression (0.0 to 1.0)
            min_box_size: Minimum box dimension in pixels
            use_vectorized: Use PyTorch implementation if available
            eps: Epsilon for numerical stability in IoU computation
        """
        if not 0.0 <= iou_threshold <= 1.0:
            raise ValueError(f"iou_threshold must be in [0, 1], got {iou_threshold}")
        
        if min_box_size < 0:
            raise ValueError(f"min_box_size must be >= 0, got {min_box_size}")
        
        self.iou_threshold = iou_threshold
        self.min_box_size = min_box_size
        self.use_vectorized = use_vectorized and TORCHVISION_AVAILABLE
        self.eps = eps
        
        self._stats = {
            "total_calls": 0,
            "vectorized_calls": 0,
            "iterative_calls": 0,
            "total_input": 0,
            "total_output": 0,
            "total_suppressed": 0
        }
    
    def apply(
        self,
        detections: List[Dict[str, Any]],
        class_agnostic: bool = True,
        target_categories: Optional[List[str]] = None
    ) -> List[Dict[str, Any]]:
        """
        Apply NMS to detections.
        
        Args:
            detections: List of detection dicts with schema:
                {
                    "category": str,
                    "confidence": float,
                    "bounding_box": {"x1": float, "y1": float, "x2": float, "y2": float},
                    ... (other fields preserved)
                }
            class_agnostic: If True, suppress across all classes
            target_categories: Optional list of categories to process (others ignored)
        
        Returns:
            Filtered list of detections with identical schema
        
        Raises:
            ValueError: If detection schema is invalid
        """
        self._stats["total_calls"] += 1
        self._stats["total_input"] += len(detections)
        
        if not detections:
            return detections
        
        if len(detections) == 1:
            self._stats["total_output"] += 1
            return detections
        
        # Validate schema
        if not all(self._validate_detection_schema(d) for d in detections):
            raise ValueError("Invalid detection schema detected")
        
        # Filter by target categories if specified
        if target_categories is not None:
            detections = [d for d in detections if d.get('category') in target_categories]
            if not detections:
                return detections
        
        # Apply NMS
        if class_agnostic:
            result = self._apply_nms_single_pass(detections)
        else:
            result = self._apply_nms_per_class(detections)
        
        self._stats["total_output"] += len(result)
        self._stats["total_suppressed"] += (len(detections) - len(result))
        
        return result
    
    def _apply_nms_single_pass(self, detections: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Apply NMS across all classes in a single pass."""
        if self.use_vectorized:
            self._stats["vectorized_calls"] += 1
            return self._nms_vectorized(detections)
        else:
            self._stats["iterative_calls"] += 1
            return self._nms_iterative(detections)
    
    def _apply_nms_per_class(self, detections: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Apply NMS separately for each class."""
        # Group by category
        category_groups = {}
        for det in detections:
            cat = det.get('category')
            if cat not in category_groups:
                category_groups[cat] = []
            category_groups[cat].append(det)
        
        # Apply NMS per category
        result = []
        for cat_dets in category_groups.values():
            if self.use_vectorized:
                result.extend(self._nms_vectorized(cat_dets))
            else:
                result.extend(self._nms_iterative(cat_dets))
        
        if self.use_vectorized:
            self._stats["vectorized_calls"] += 1
        else:
            self._stats["iterative_calls"] += 1
        
        return result
    
    def _nms_vectorized(self, detections: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """
        Vectorized NMS using torchvision.ops.nms.
        
        This matches YOLO's built-in agnostic NMS exactly.
        """
        if not detections:
            return []
        
        # Filter invalid boxes
        valid_dets = [d for d in detections if self._validate_box(d["bounding_box"])]
        
        if not valid_dets:
            return []
        
        try:
            # Convert to tensors
            boxes = torch.tensor([
                [d["bounding_box"]["x1"], d["bounding_box"]["y1"],
                 d["bounding_box"]["x2"], d["bounding_box"]["y2"]]
                for d in valid_dets
            ], dtype=torch.float32)
            
            scores = torch.tensor([d["confidence"] for d in valid_dets], dtype=torch.float32)
            
            # Apply torchvision NMS
            keep_indices = torchvision_nms(boxes, scores, self.iou_threshold)
            keep_indices = keep_indices.cpu().numpy()
            
            return [valid_dets[i] for i in keep_indices]
        
        except Exception as e:
            # Fallback to iterative on error
            return self._nms_iterative(valid_dets)
    
    def _nms_iterative(self, detections: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """
        Iterative NMS implementation with YOLO-matching enhancements.
        
        This provides equivalent results to vectorized NMS without PyTorch.
        """
        if not detections:
            return []
        
        # Filter invalid boxes
        valid_dets = [d for d in detections if self._validate_box(d["bounding_box"])]
        
        if not valid_dets:
            return []
        
        # Sort by confidence with area tie-breaking
        def sort_key(d):
            bbox = d["bounding_box"]
            area = abs(bbox['x2'] - bbox['x1']) * abs(bbox['y2'] - bbox['y1'])
            return (d["confidence"], area)
        
        sorted_dets = sorted(valid_dets, key=sort_key, reverse=True)
        
        # Apply NMS
        keep = []
        suppressed_indices = set()
        
        for i, det in enumerate(sorted_dets):
            if i in suppressed_indices:
                continue
            
            keep.append(det)
            best_bbox = det["bounding_box"]
            
            # Suppress overlapping boxes
            for j in range(i + 1, len(sorted_dets)):
                if j in suppressed_indices:
                    continue
                
                other_bbox = sorted_dets[j]["bounding_box"]
                iou = self._compute_iou(best_bbox, other_bbox)
                
                # Use >= for consistency with torchvision
                if iou >= self.iou_threshold:
                    suppressed_indices.add(j)
        
        return keep
    
    def _compute_iou(self, bbox1: Dict, bbox2: Dict) -> float:
        """
        Compute IoU with numerical stability.
        
        Args:
            bbox1: First box {x1, y1, x2, y2}
            bbox2: Second box {x1, y1, x2, y2}
        
        Returns:
            IoU value in [0, 1]
        """
        # Ensure coordinates are in correct order
        x1_1 = min(bbox1['x1'], bbox1['x2'])
        y1_1 = min(bbox1['y1'], bbox1['y2'])
        x2_1 = max(bbox1['x1'], bbox1['x2'])
        y2_1 = max(bbox1['y1'], bbox1['y2'])
        
        x1_2 = min(bbox2['x1'], bbox2['x2'])
        y1_2 = min(bbox2['y1'], bbox2['y2'])
        x2_2 = max(bbox2['x1'], bbox2['x2'])
        y2_2 = max(bbox2['y1'], bbox2['y2'])
        
        # Compute intersection
        inter_x1 = max(x1_1, x1_2)
        inter_y1 = max(y1_1, y1_2)
        inter_x2 = min(x2_1, x2_2)
        inter_y2 = min(y2_1, y2_2)
        
        inter_w = max(0.0, inter_x2 - inter_x1)
        inter_h = max(0.0, inter_y2 - inter_y1)
        
        if inter_w == 0.0 or inter_h == 0.0:
            return 0.0
        
        inter_area = inter_w * inter_h
        
        # Compute box areas
        area1 = (x2_1 - x1_1) * (y2_1 - y1_1)
        area2 = (x2_2 - x1_2) * (y2_2 - y1_2)
        
        # Add epsilon for numerical stability (matches torchvision)
        union_area = area1 + area2 - inter_area + self.eps
        
        # Compute IoU with safeguards
        iou = inter_area / union_area
        
        # Clamp to valid range
        return max(0.0, min(1.0, iou))
    
    def _validate_box(self, bbox: Dict, max_wh: float = 1e4) -> bool:
        """
        Validate box dimensions (matches YOLO's internal validation).
        
        Args:
            bbox: Box dict {x1, y1, x2, y2}
            max_wh: Maximum box width/height
        
        Returns:
            True if box is valid
        """
        try:
            w = abs(bbox['x2'] - bbox['x1'])
            h = abs(bbox['y2'] - bbox['y1'])
            
            # Check minimum size
            if w < self.min_box_size or h < self.min_box_size:
                return False
            
            # Check maximum size
            if w > max_wh or h > max_wh:
                return False
            
            # Check for valid numeric values
            for v in [bbox['x1'], bbox['y1'], bbox['x2'], bbox['y2']]:
                if not isinstance(v, (int, float)):
                    return False
                if np.isnan(v) or np.isinf(v):
                    return False
            
            return True
        
        except (KeyError, TypeError, ValueError):
            return False
    
    def _validate_detection_schema(self, detection: Dict[str, Any]) -> bool:
        """
        Validate detection has required fields for NMS.
        
        Args:
            detection: Single detection dict
        
        Returns:
            True if schema is valid
        """
        if not isinstance(detection, dict):
            return False
        
        # Check category
        if 'category' not in detection or not isinstance(detection['category'], str):
            return False
        
        # Check confidence
        if 'confidence' not in detection:
            return False
        conf = detection['confidence']
        if not isinstance(conf, (int, float)) or conf < 0 or conf > 1:
            return False
        
        # Check bounding_box
        bbox = detection.get('bounding_box')
        if not isinstance(bbox, dict):
            return False
        
        required_keys = {'x1', 'y1', 'x2', 'y2'}
        if not required_keys.issubset(bbox.keys()):
            return False
        
        # Validate bbox coordinates are numeric
        try:
            for v in [bbox['x1'], bbox['y1'], bbox['x2'], bbox['y2']]:
                if not isinstance(v, (int, float)):
                    return False
        except (KeyError, TypeError):
            return False
        
        return True
    
    def get_stats(self) -> Dict[str, Any]:
        """
        Get NMS usage statistics.
        
        Returns:
            Dictionary with statistics:
            - total_calls: Number of times apply() was called
            - vectorized_calls: Number of vectorized NMS calls
            - iterative_calls: Number of iterative NMS calls
            - total_input: Total input detections
            - total_output: Total output detections
            - total_suppressed: Total suppressed detections
            - suppression_rate: Percentage of detections suppressed
        """
        stats = self._stats.copy()
        if stats["total_input"] > 0:
            stats["suppression_rate"] = 100 * stats["total_suppressed"] / stats["total_input"]
        else:
            stats["suppression_rate"] = 0.0
        return stats
    
    def reset_stats(self):
        """Reset usage statistics."""
        self._stats = {
            "total_calls": 0,
            "vectorized_calls": 0,
            "iterative_calls": 0,
            "total_input": 0,
            "total_output": 0,
            "total_suppressed": 0
        }
    
    @staticmethod
    def is_vectorized_available() -> bool:
        """Check if vectorized implementation is available."""
        return TORCHVISION_AVAILABLE


# Convenience function for quick usage
def apply_nms(
    detections: List[Dict[str, Any]],
    iou_threshold: float = 0.45,
    class_agnostic: bool = True,
    min_box_size: float = 2.0,
    use_vectorized: bool = True
) -> List[Dict[str, Any]]:
    """
    Convenience function for one-time NMS application.
    
    Args:
        detections: List of detection dicts
        iou_threshold: IoU threshold for suppression
        class_agnostic: If True, suppress across all classes
        min_box_size: Minimum box dimension in pixels
        use_vectorized: Use PyTorch implementation if available
    
    Returns:
        Filtered list of detections
    
    Example:
        >>> detections = [
        ...     {"category": "car", "confidence": 0.9, 
        ...      "bounding_box": {"x1": 100, "y1": 100, "x2": 200, "y2": 200}},
        ...     {"category": "car", "confidence": 0.85,
        ...      "bounding_box": {"x1": 105, "y1": 105, "x2": 205, "y2": 205}}
        ... ]
        >>> filtered = apply_nms(detections, iou_threshold=0.5, class_agnostic=True)
        >>> len(filtered)
        1
    """
    nms = AgnosticNMS(
        iou_threshold=iou_threshold,
        min_box_size=min_box_size,
        use_vectorized=use_vectorized
    )
    return nms.apply(detections, class_agnostic=class_agnostic)