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 = 250.0,
            sbend_penalty: float = 500.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.
            sbend_penalty: Base cost for parametric S-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,
            sbend_penalty=sbend_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 + mandatory turn penalties.
        """
        dx = abs(current.x - target.x)
        dy = abs(current.y - target.y)
        dist = dx + dy

        penalty = 0.0
        if abs(current.orientation - target.orientation) > 0.1:
            # Needs at least 1 bend
            penalty += 10.0 + self.config.bend_penalty * 0.1

        return self.greedy_h_weight * (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,
            skip_congestion: bool = False,
            penalty: float = 0.0,
            ) -> 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.
            skip_congestion: If True, bypass congestion checks.
            penalty: Fixed cost penalty for the move type.

        Returns:
            Total cost of the move, or 1e15 if invalid.
        """
        _ = net_width  # Unused
        
        # 1. Boundary Check
        danger_map = self.danger_map
        if not danger_map.is_within_bounds(end_port.x, end_port.y):
            return 1e15

        total_cost = length * self.unit_length_cost + penalty

        # 2. Collision Check
        # FAST PATH: skip_static and skip_congestion are often True when called from optimized AStar
        if not skip_static or not skip_congestion:
            collision_engine = self.collision_engine
            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 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)
                if not skip_congestion:
                    overlaps = 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 += danger_map.get_cost(end_port.x, end_port.y)
        return total_cost
