pathfinder.py

import heapq
from typing import Optional
from schemas import MapDataDTO, ZoneDTO


class ReservationTable:
    """Centralized space-time matrix tracking dynamic constraints"""

    def __init__(self) -> None:
        """Initializes the class object"""
        # Maps (node_name, time_step), current_occupancy_count
        self.node_res: dict[tuple[str, int], int] = {}
        # Maps (edge_name, time_step), current_occupancy_count
        self.edge_res: dict[tuple[str, int], int] = {}

    def reserve_node(self, node: str, t: int) -> None:
        """Books a capacity slot inside a node for a turn"""
        self.node_res[(node, t)] = self.node_res.get((node, t), 0) + 1

    def reserve_edge(self, edge: str, t: int) -> None:
        """Books a bidirectional edge corridor for a turn"""
        self.edge_res[(edge, t)] = self.edge_res.get((edge, t), 0) + 1

    def get_node_occupancy(self, node: str, t: int) -> int:
        """Returns current scheduled drone count inside a node at T"""
        return self.node_res.get((node, t), 0)

    def get_edge_occupancy(self, edge: str, t: int) -> int:
        """Returns current scheduled transit volume on edge at T"""
        return self.edge_res.get((edge, t), 0)


class SpaceTimeNode:
    """Pathfinding search state in space and time"""

    def __init__(
            self,
            node_name: str,
            t: int,
            g_cost: int,
            h_cost: float,
            parent: Optional["SpaceTimeNode"] = None,
            action_str: str = ""
    ) -> None:
        """Initializes the class object"""
        self.node_name: str = node_name
        self.t: int = t
        # actual steps taken so far - guidance cost
        self.g_cost: int = g_cost
        # estimated steps to goal (heuristic cost sets priority)
        self.h_cost: float = h_cost
        self.f_cost: float = float(g_cost) + h_cost
        self.parent: Optional[SpaceTimeNode] = parent
        # movement ("D1-gate1")
        self.action_str: str = action_str

    def __lt__(self, other: "SpaceTimeNode") -> bool:
        """Heuristic favoring lower cost states in heapq"""
        if self.f_cost == other.f_cost:
            return self.g_cost > other.g_cost
        return self.f_cost < other.f_cost


class SpaceTimePathfinder:
    """Executes cooperative space-time routing across constrained nodes"""

    def __init__(self,
                 map_data: MapDataDTO,
                 ground_mode: bool = False) -> None:
        self.map_data: MapDataDTO = map_data
        self.ground_mode: bool = ground_mode
        self.res_table: ReservationTable = ReservationTable()
        self.heuristics: dict[str, float] = {}

        # Sets fast look-up table
        self.all_nodes: dict[str, ZoneDTO] = {
            map_data.start_hub.name: map_data.start_hub,
            map_data.end_hub.name: map_data.end_hub
        }
        # Fills the dict with hubs from map_data
        # [hub_name, ZoneDTO(name='start', x=0, y=0,
        # zone_type='normal', max_drones]
        self.all_nodes.update(map_data.hubs)

        # Sets empty runtime adjacency maps
        self.adj_map: dict[str, list[tuple[str, str, int, str]]] = {
            name: [] for name in self.all_nodes
        }

        self._build_graph()
        self._compute_heuristics()

    def _build_graph(self) -> None:
        """Maps bidirectional corridors to adjacency structures"""

        for conn_name, conn_dto in self.map_data.connections.items():
            z1, z2 = conn_dto.zone1, conn_dto.zone2
            cap: int = conn_dto.max_link_capacity

            if self.all_nodes[z1].zone_type == "blocked" or \
               self.all_nodes[z2].zone_type == "blocked":
                continue

            ztype1: str = self.all_nodes[z1].zone_type
            ztype2: str = self.all_nodes[z2].zone_type

            self.adj_map[z1].append((z2, conn_name, cap, ztype2))
            self.adj_map[z2].append((z1, conn_name, cap, ztype1))
            # adj. map:
            # {'start': [('waypoint1', 'start-waypoint1', 1,
            # 'normal')],
            #  'goal': [('waypoint2', 'waypoint2-goal', 1,
            # 'normal')],
            #  'waypoint1': [('start', 'start-waypoint1', 1,
            # 'normal'), ('waypoint2', 'waypoint1-waypoint2', 1,
            # 'normal')],
            # 'waypoint2': [('waypoint1', 'waypoint1-waypoint2',
            #  1, 'normal'), ('goal', 'waypoint2-goal', 1, 'normal')]}

    def _compute_heuristics(self) -> None:
        """Reverse Dijkstra from goal pre-calculating guidance cost"""

        goal_name: str = self.map_data.end_hub.name
        self.heuristics = {name: 999999.0 for name in self.all_nodes}
        self.heuristics[goal_name] = 0.0

        pq: list[tuple[float, str]] = [(0.0, goal_name)]

        while pq:
            current_cost, current_node = heapq.heappop(pq)
            if current_cost > self.heuristics[current_node]:
                continue

            for neighbor, _, _, n_type in self.adj_map[current_node]:
                step_cost: float = 2.0 if n_type == "restricted" else 1.0

                if n_type == "priority":
                    step_cost -= 0.1

                new_cost: float = current_cost + step_cost

                if new_cost < self.heuristics[neighbor]:
                    self.heuristics[neighbor] = new_cost
                    heapq.heappush(pq, (new_cost, neighbor))
        # self.heuristics = {goal: 0, waypoint2: 1, waypoint1: 3, start: 4}

    def route_fleet(self) -> dict[int, list[SpaceTimeNode]]:
        """Schedules space-time paths for all fleet units"""

        fleet_routes: dict[int, list[SpaceTimeNode]] = {}
        start_name: str = self.map_data.start_hub.name
        goal_name: str = self.map_data.end_hub.name

        for drone_id in range(1, self.map_data.nb_drones + 1):
            path = self._find_path_for_drone(drone_id, start_name, goal_name)

            if not path:
                raise RuntimeError("CRITICAL DEADLOCK: Unsolvable path "
                                   f"for Drone {drone_id}.")

            fleet_routes[drone_id] = path

            for node in path:
                if node.t > 0:
                    if node.node_name != goal_name:
                        self.res_table.reserve_node(node.node_name, node.t)
                    if "-" in node.action_str:
                        parts = node.action_str.split("-")
                        if len(parts) == 3:
                            edge_str = f"{parts[1]}-{parts[2]}"
                            self.res_table.reserve_edge(edge_str, node.t - 1)

        return fleet_routes

    def _find_path_for_drone(
            self,
            drone_id: int,
            start_node: str,
            goal_node: str
    ) -> list[SpaceTimeNode]:
        """Core space-time A* search with collision look-aheads"""

        open_set: list[SpaceTimeNode] = []
        closed_set: set[tuple[str, int]] = set()

        start_state = SpaceTimeNode(
            node_name=start_node,
            t=0,
            g_cost=0,
            h_cost=self.heuristics.get(start_node, 0.0),
            action_str=f"D{drone_id}-{start_node}"
        )
        heapq.heappush(open_set, start_state)

        while open_set:
            current = heapq.heappop(open_set)
            state_key = (current.node_name, current.t)

            if current.t > 200:
                continue

            if state_key in closed_set:
                continue
            closed_set.add(state_key)

            if current.node_name == goal_node:
                return self._reconstruct_path(current)

            edges = self.adj_map[current.node_name]
            for neighbor, edge_name, link_cap, n_type in edges:
                step_cost = 2 if n_type == "restricted" else 1
                next_t = current.t + step_cost

                edge_occ = self.res_table.get_edge_occupancy(
                    edge_name, current.t)
                if edge_occ >= link_cap:
                    continue

                target_node_cap = self.all_nodes[neighbor].max_drones
                if neighbor != goal_node:
                    node_occ = self.res_table.get_node_occupancy(
                        neighbor, next_t)
                    if node_occ >= target_node_cap:
                        continue

                if n_type == "restricted":
                    move_str = f"D{drone_id}-{edge_name}"
                else:
                    move_str = f"D{drone_id}-{neighbor}"

                neighbor_state = SpaceTimeNode(
                    node_name=neighbor,
                    t=next_t,
                    g_cost=current.g_cost + step_cost,
                    h_cost=self.heuristics.get(neighbor, 0.0),
                    parent=current,
                    action_str=move_str
                )
                heapq.heappush(open_set, neighbor_state)

            current_cap = self.all_nodes[current.node_name].max_drones
            next_occ = self.res_table.get_node_occupancy(
                current.node_name, current.t + 1)

            if current.node_name == start_node or next_occ < current_cap:
                wait_state = SpaceTimeNode(
                    node_name=current.node_name,
                    t=current.t + 1,
                    g_cost=current.g_cost + 1,
                    h_cost=current.h_cost,
                    parent=current,
                    action_str=f"D{drone_id}-{current.node_name}"
                )
                heapq.heappush(open_set, wait_state)

        return []

    def _reconstruct_path(
            self, end_node: SpaceTimeNode
    ) -> list[SpaceTimeNode]:
        """Backtracks parent pointers to output clean step chains."""
        path: list[SpaceTimeNode] = []
        curr: Optional[SpaceTimeNode] = end_node
        while curr:
            path.append(curr)
            curr = curr.parent
        return path[::-1]