simulation.py

import heapq
from module import Graph, Zone, Connection, Drone, EndHub
from module import Visualization as Visual
from parser import MapData


class UnsolvableGraphError(Exception):
    """Exception raised when a graph algorithm cannot find a solution."""
    pass


class Engine:
    def __init__(self, data: MapData, graph: Graph) -> None:
        """Create a simulation engine.

        Args:
            data: Parsed map data.
            graph: Graph built from the parsed map.
        """
        self.data = data
        self.graph = graph
        self.node_obj = {
            node.name: node for node in graph.adj_list.keys()
        }

    def dijkstra(self, start_zone: Zone) -> list[str]:
        """Compute the path from the start zone to the final zone.

        Args:
            start_zone: Zone where the search begins.

        Returns:
            The ordered list of zone names for the chosen path.
        """
        distances: dict[str, float] = {
          node.name: float('inf') for node in self.graph.adj_list.keys()
        }
        distances[start_zone.name] = 0

        predecessors: dict[str, None | str] = {
            node.name: None for node in self.graph.adj_list.keys()
        }

        heap: list[tuple[int | float, int, str]] = [
            (0, start_zone.priority, start_zone.name)
        ]

        visited = []
        while heap:
            zone_cost, zone_priority, zone_name = heapq.heappop(heap)

            zone_node = self.node_obj[zone_name]
            for neighbor_cost, neighbor in self.graph.get_neighbors(zone_node):

                if (neighbor.name not in visited
                   and neighbor_cost + zone_cost < distances[neighbor.name]):

                    distances[neighbor.name] = neighbor_cost + zone_cost
                    predecessors[neighbor.name] = zone_name

                    heapq.heappush(
                        heap, (distances[neighbor.name],
                               neighbor.priority + zone_priority,
                               neighbor.name)
                    )

            visited.append(zone_name)

        return self.get_path(predecessors)

    def get_path(self, predecessors: dict[str, None | str]) -> list[str]:
        """Reconstruct a path from predecessor links.

        Args:
            predecessors: Mapping of node name to predecessor name.

        Returns:
            The reconstructed path as a list of zone names.

        Raises:
            UnsolvableGraphError: If the path cannot reach a destination.
        """
        path = []

        node: str | None = next(reversed(predecessors.keys()))

        while node is not None:
            path.append(node)
            node = predecessors[node]

        path.reverse()

        if len(path) == 1:
            raise UnsolvableGraphError(
                "The graph cannot be solved: no valid path or solution exists"
            )

        return path

    def get_connection(self, zone1: Zone, zone2: Zone) -> Connection:
        """Return the connection between two zones.

        Args:
            zone1: First endpoint.
            zone2: Second endpoint.

        Returns:
            The matching connection.

        Raises:
            ValueError: If no connection exists between the zones.
        """
        for connection in self.data.connection_list:
            if connection.zone1 == zone1 and connection.zone2 == zone2:
                return connection
            elif connection.zone1 == zone2 and connection.zone2 == zone1:
                return connection
        raise ValueError("No connection")

    def simulation(self, drones: list[Drone]) -> None:
        """Run the turn-based drone simulation.

        Args:
            drones: Active drones to advance through the graph.
        """
        console = Visual.Console()

        while drones:
            connection_list = []

            for drone in drones[:]:
                drone.next_zone = self.node_obj[drone.path.pop(0)]
                zone_connection = self.get_connection(
                    drone.position, drone.next_zone
                )
                connection_list.append(zone_connection)

                if drone.next_zone.default_metadata['zone'] == "restricted":
                    drone.transit_turns_left += 1

                if drone.transit_turns_left == 2:
                    drone.next_zone.available += 1

                if (drone.next_zone.available and
                   zone_connection.default_metadata["max_link_capacity"]):

                    drone.next_zone.available -= 1
                    zone_connection.default_metadata["max_link_capacity"] -= 1

                    if drone.transit_turns_left == 1:
                        next_zone_name = Visual.apply_color(drone.next_zone)
                        pre_zone_name = Visual.apply_color(drone.position)

                        console.print(
                            f"{drone.id}-"
                            f"{pre_zone_name}-"
                            f"{next_zone_name}",
                            end=" "
                        )

                        drone.path.insert(0, drone.next_zone.name)
                    else:
                        drone.position.available += 1
                        drone.position = drone.next_zone

                        drone.transit_turns_left = 0

                        position_name = Visual.apply_color(drone.position)
                        console.print(
                            f"{drone.id}-"
                            f"{position_name}", end=" "
                        )

                else:
                    drone.transit_turns_left = 0
                    drone.path.insert(0, drone.next_zone.name)

                if isinstance(drone.position, EndHub):
                    drones.remove(drone)

            for connection in connection_list:
                connection.default_metadata[
                    "max_link_capacity"
                ] = connection.fallback_max_link_capacity

            print()