mice/engine/collision_system.py

"""
Native Python collision detection system using Spatial Hashing.

This module provides efficient collision detection without NumPy.
It uses a grid-based approach (buckets) to ensure O(1) or O(n) complexity.
This structure is designed to be easily portable to Nim.
"""

from typing import Dict, List, Tuple, Set
from dataclasses import dataclass


def shrink_bbox(bbox: Tuple[float, float, float, float], ratio: float) -> Tuple[float, float, float, float]:
    """Return a smaller bbox centered on the original one.

    ratio is the fraction of width/height removed from each side.
    E.g. ratio=0.25 keeps the central 50% of the original area.
    """
    x1, y1, x2, y2 = bbox
    dx = (x2 - x1) * ratio / 2
    dy = (y2 - y1) * ratio / 2
    return (x1 + dx, y1 + dy, x2 - dx, y2 - dy)


@dataclass
class CollisionLayer:
    """Define which types of units can collide with each other."""
    RAT = 0
    BOMB = 1
    GAS = 2
    MINE = 3
    POINT = 4
    EXPLOSION = 5


class CollisionSystem:
    """
    Manages collision detection using a Spatial Grid.
    """
    
    def __init__(self, cell_size: int, grid_width: int, grid_height: int):
        self.cell_size = cell_size
        self.grid_width = grid_width
        self.grid_height = grid_height
        
        # Grid: Maps (x, y) coordinates to list of Unit objects/IDs
        self.grid: Dict[Tuple[int, int], List[int]] = {}
        self.grid_before: Dict[Tuple[int, int], List[int]] = {}
        
        # Unit storage
        self.units_data: Dict[int, dict] = {}
        self.unit_ids: List[int] = [] # Stable list of IDs for parity
        
        # Collision matrix (Native Python dict of sets for speed)
        self._setup_collision_matrix()
        
    def _setup_collision_matrix(self):
        """Define which collision layers interact with each other."""
        L = CollisionLayer
        # Interaction rules: layer -> set of target layers
        self.interaction_map = {
            L.RAT: {L.RAT, L.GAS, L.MINE, L.POINT, L.EXPLOSION},
            L.GAS: {L.RAT},
            L.MINE: {L.RAT},
            L.POINT: {L.RAT},
            L.EXPLOSION: {L.RAT},
            L.BOMB: set() # Bombs are passive until they explode
        }
        
    def clear(self):
        """Clear all collision data for new frame."""
        self.grid.clear()
        self.grid_before.clear()
        self.units_data.clear()
        self.unit_ids.clear()
    
    def register_unit(self, unit_id, bbox: Tuple[float, float, float, float], 
                     position: Tuple[int, int], position_before: Tuple[int, int],
                     layer: int):
        """
        Register a unit in the spatial grid.
        """
        self.unit_ids.append(unit_id)
        self.units_data[unit_id] = {
            "bbox": bbox,
            "pos": position,
            "pos_before": position_before,
            "layer": layer
        }
        
        # Add to spatial buckets
        if position not in self.grid:
            self.grid[position] = []
        self.grid[position].append(unit_id)
        
        if position_before not in self.grid_before:
            self.grid_before[position_before] = []
        self.grid_before[position_before].append(unit_id)
    
    def get_collisions_for_unit(self, unit_id, layer: int, 
                               tolerance: int = 0) -> List[Tuple[int, any]]:
        """
        Get all units colliding with the specified unit using grid lookup.
        """
        if unit_id not in self.units_data:
            return []
            
        data = self.units_data[unit_id]
        bbox = data["bbox"]
        pos = data["pos"]
        pos_before = data["pos_before"]
        
        colliding_units = []
        target_layers = self.interaction_map.get(layer, set())
        
        # Candidate search: look in current and previous grid buckets
        # This covers units that moved into our space or were there before
        candidates = set()
        for p in [pos, pos_before]:
            if p in self.grid:
                candidates.update(self.grid[p])
            if p in self.grid_before:
                candidates.update(self.grid_before[p])
        
        candidates.discard(unit_id)
        
        for other_id in candidates:
            other_data = self.units_data[other_id]
            
            # 1. Filter by layer
            if other_data["layer"] not in target_layers:
                continue
                
            # 2. AABB Check
            other_bbox = other_data["bbox"]
            if (bbox[0] < other_bbox[2] - tolerance and
                bbox[2] > other_bbox[0] + tolerance and
                bbox[1] < other_bbox[3] - tolerance and
                bbox[3] > other_bbox[1] + tolerance):
                # Return dummy index (for parity) and ID
                colliding_units.append((0, other_id))
                
        return colliding_units
    
    def get_units_in_cell(self, position: Tuple[int, int], 
                         use_before: bool = False) -> List[any]:
        """Get all unit IDs in a specific grid cell."""
        target_grid = self.grid_before if use_before else self.grid
        return target_grid.get(position, [])
    
    def get_units_in_area(self, positions: List[Tuple[int, int]], 
                         layer_filter: int = None) -> Set[any]:
        """Get all units in multiple grid cells (vectorized lookup replacement)."""
        found = set()
        for pos in positions:
            # Check current grid
            if pos in self.grid:
                for uid in self.grid[pos]:
                    if layer_filter is None or self.units_data[uid]["layer"] == layer_filter:
                        found.add(uid)
            # Check previous grid
            if pos in self.grid_before:
                for uid in self.grid_before[pos]:
                    if layer_filter is None or self.units_data[uid]["layer"] == layer_filter:
                        found.add(uid)
        return found

    def check_partial_move_collision(self, unit_id, partial_move: float, 
                                    threshold: float = 0.5) -> List[any]:
        """Collision check considering movement progress."""
        if unit_id not in self.units_data:
            return []
            
        data = self.units_data[unit_id]
        pos = data["pos"] if partial_move >= threshold else data["pos_before"]
        
        found = set()
        if pos in self.grid:
            found.update(self.grid[pos])
        if pos in self.grid_before:
            found.update(self.grid_before[pos])
            
        found.discard(unit_id)
        return list(found)