from __future__ import annotations

from typing import TYPE_CHECKING
import numpy
import shapely

if TYPE_CHECKING:
    from shapely.geometry import Polygon


class DangerMap:
    """
    A pre-computed grid for heuristic proximity costs, vectorized for performance.
    """
    __slots__ = ('minx', 'miny', 'maxx', 'maxy', 'resolution', 'safety_threshold', 'k', 'width_cells', 'height_cells', 'grid')

    minx: float
    miny: float
    maxx: float
    maxy: float
    """ Boundary coordinates of the map """

    resolution: float
    """ Grid cell size in micrometers """

    safety_threshold: float
    """ Distance below which proximity costs are applied """

    k: float
    """ Cost multiplier constant """

    width_cells: int
    height_cells: int
    """ Grid dimensions in cells """

    grid: numpy.ndarray
    """ 2D array of pre-computed costs """

    def __init__(
            self,
            bounds: tuple[float, float, float, float],
            resolution: float = 1.0,
            safety_threshold: float = 10.0,
            k: float = 1.0,
            ) -> None:
        """
        Initialize the Danger Map.

        Args:
            bounds: (minx, miny, maxx, maxy) in um.
            resolution: Cell size (um).
            safety_threshold: Proximity limit (um).
            k: Penalty multiplier.
        """
        self.minx, self.miny, self.maxx, self.maxy = bounds
        self.resolution = resolution
        self.safety_threshold = safety_threshold
        self.k = k

        # Grid dimensions
        self.width_cells = int(numpy.ceil((self.maxx - self.minx) / self.resolution))
        self.height_cells = int(numpy.ceil((self.maxy - self.miny) / self.resolution))

        self.grid = numpy.zeros((self.width_cells, self.height_cells), dtype=numpy.float32)

    def precompute(self, obstacles: list[Polygon]) -> None:
        """
        Pre-compute the proximity costs for the entire grid using vectorized operations.

        Args:
            obstacles: List of static obstacle geometries.
        """
        from scipy.ndimage import distance_transform_edt

        # 1. Create a binary mask of obstacles
        mask = numpy.ones((self.width_cells, self.height_cells), dtype=bool)

        # Create coordinate grids
        x_coords = numpy.linspace(self.minx + self.resolution/2, self.maxx - self.resolution/2, self.width_cells)
        y_coords = numpy.linspace(self.miny + self.resolution/2, self.maxy - self.resolution/2, self.height_cells)
        xv, yv = numpy.meshgrid(x_coords, y_coords, indexing='ij')

        for poly in obstacles:
            # Use shapely.contains_xy for fast vectorized point-in-polygon check
            in_poly = shapely.contains_xy(poly, xv, yv)
            mask[in_poly] = False

        # 2. Distance transform (mask=True for empty space)
        distances = distance_transform_edt(mask) * self.resolution

        # 3. Proximity cost: k / d^2 if d < threshold, else 0
        # Cap distances at a small epsilon (e.g. 0.1um) to avoid division by zero
        safe_distances = numpy.maximum(distances, 0.1)
        self.grid = numpy.where(
            distances < self.safety_threshold,
            self.k / (safe_distances**2),
            0.0
        ).astype(numpy.float32)

    def is_within_bounds(self, x: float, y: float) -> bool:
        """
        Check if a coordinate is within the design bounds.

        Args:
            x, y: Coordinate to check.

        Returns:
            True if within [min, max] for both axes.
        """
        return self.minx <= x <= self.maxx and self.miny <= y <= self.maxy

    def get_cost(self, x: float, y: float) -> float:
        """
        Get the proximity cost at a specific coordinate.

        Args:
            x, y: Coordinate to look up.

        Returns:
            Pre-computed cost, or 1e15 if out of bounds.
        """
        ix = int((x - self.minx) / self.resolution)
        iy = int((y - self.miny) / self.resolution)

        if 0 <= ix < self.width_cells and 0 <= iy < self.height_cells:
            return float(self.grid[ix, iy])
        return 1e15  # Outside bounds is impossible
