inire/inire/router/cost.py

from __future__ import annotations

from typing import TYPE_CHECKING

from inire.router.config import CostConfig

if TYPE_CHECKING:
    from shapely.geometry import Polygon

    from inire.geometry.collision import CollisionEngine
    from inire.geometry.primitives import Port
    from inire.router.danger_map import DangerMap


class CostEvaluator:
    """
    Calculates total path and proximity costs.
    """
    __slots__ = ('collision_engine', 'danger_map', 'config', 'unit_length_cost', 'greedy_h_weight', 'congestion_penalty')

    collision_engine: CollisionEngine
    """ The engine for intersection checks """

    danger_map: DangerMap
    """ Pre-computed grid for heuristic proximity costs """

    config: CostConfig
    """ Parameter configuration """

    unit_length_cost: float
    greedy_h_weight: float
    congestion_penalty: float
    """ Cached weight values for performance """

    def __init__(
            self,
            collision_engine: CollisionEngine,
            danger_map: DangerMap,
            unit_length_cost: float = 1.0,
            greedy_h_weight: float = 1.1,
            congestion_penalty: float = 10000.0,
            bend_penalty: float = 50.0,
            ) -> None:
        """
        Initialize the Cost Evaluator.

        Args:
            collision_engine: The engine for intersection checks.
            danger_map: Pre-computed grid for heuristic proximity costs.
            unit_length_cost: Cost multiplier per micrometer of path length.
            greedy_h_weight: Heuristic weighting (A* greedy factor).
            congestion_penalty: Multiplier for path overlaps in negotiated congestion.
            bend_penalty: Base cost for 90-degree bends.
        """
        self.collision_engine = collision_engine
        self.danger_map = danger_map
        self.config = CostConfig(
            unit_length_cost=unit_length_cost,
            greedy_h_weight=greedy_h_weight,
            congestion_penalty=congestion_penalty,
            bend_penalty=bend_penalty,
        )

        # Use config values
        self.unit_length_cost = self.config.unit_length_cost
        self.greedy_h_weight = self.config.greedy_h_weight
        self.congestion_penalty = self.config.congestion_penalty


    def g_proximity(self, x: float, y: float) -> float:
        """
        Get proximity cost from the Danger Map.

        Args:
            x, y: Coordinate to check.

        Returns:
            Proximity cost at location.
        """
        return self.danger_map.get_cost(x, y)

    def h_manhattan(self, current: Port, target: Port) -> float:
        """
        Heuristic: weighted Manhattan distance + orientation penalty.

        Args:
            current: Current port state.
            target: Target port state.

        Returns:
            Heuristic cost estimate.
        """
        dx = abs(current.x - target.x)
        dy = abs(current.y - target.y)
        dist = dx + dy

        # Orientation penalty if not aligned with target entry
        # If we need to turn, the cost is at least min_bend_radius * pi/2
        # But we also need to account for the physical distance required for the turn.
        penalty = 0.0
        if current.orientation != target.orientation:
            # 90-degree turn cost: radius 10 -> ~15.7 um + penalty
            penalty += 15.7 + self.config.bend_penalty

        # Add 1.5 multiplier for greediness (faster search)
        return 1.5 * (dist + penalty)


    def evaluate_move(
            self,
            geometry: list[Polygon],
            end_port: Port,
            net_width: float,
            net_id: str,
            start_port: Port | None = None,
            length: float = 0.0,
            dilated_geometry: list[Polygon] | None = None,
            skip_static: bool = False,
            ) -> float:
        """
        Calculate the cost of a single move (Straight, Bend, SBend).

        Args:
            geometry: List of polygons in the move.
            end_port: Port at the end of the move.
            net_width: Width of the waveguide (unused).
            net_id: Identifier for the net.
            start_port: Port at the start of the move.
            length: Physical path length of the move.
            dilated_geometry: Pre-calculated dilated polygons.
            skip_static: If True, bypass static collision checks (e.g. if already done).

        Returns:
            Total cost of the move, or 1e15 if invalid.
        """
        _ = net_width  # Unused
        total_cost = length * self.unit_length_cost

        # 1. Boundary Check
        if not self.danger_map.is_within_bounds(end_port.x, end_port.y):
            return 1e15

        # 2. Collision Check
        for i, poly in enumerate(geometry):
            dil_poly = dilated_geometry[i] if dilated_geometry else None
            # Hard Collision (Static obstacles)
            if not skip_static:
                if self.collision_engine.check_collision(
                    poly, net_id, buffer_mode='static', start_port=start_port, end_port=end_port,
                    dilated_geometry=dil_poly
                ):
                    return 1e15

            # Soft Collision (Negotiated Congestion)
            overlaps = self.collision_engine.check_collision(
                poly, net_id, buffer_mode='congestion', dilated_geometry=dil_poly
            )
            if isinstance(overlaps, int) and overlaps > 0:
                total_cost += overlaps * self.congestion_penalty

        # 3. Proximity cost from Danger Map
        total_cost += self.g_proximity(end_port.x, end_port.y)
        return total_cost