from math import inf
from typing import Iterator, Optional

from src.graph.node import Node
from src.graph.edge import Edge

class Graph:
    def __init__(self):
        self.edges: list[Edge] = []
        self.nodes: list[Node] = []

    def add_node(self, x: int, y: int, name: str) -> None:
        self.nodes.append(Node(x, y, name))

    def add_edge(self, start_index: int, end_index: int, length: float) -> None:
        self.edges.append(Edge(start_index, end_index, length))

    def edges_adjacent_to(self, node_i: int) -> Iterator[Edge]:
        return filter(lambda e: e.start == node_i or e.end == node_i, self.edges)

    def dijkstra(self, source_index: int, target_index: int) -> Optional[list[int]]:
        n = len(self.nodes)

        if source_index < 0 or source_index >= n:
            return None

        if target_index < 0 or target_index >= n:
            return None

        unvisited = list(range(n))

        distances_from_start = [inf] * n
        distances_from_start[source_index] = 0

        node_sequences = [[] for _ in range(n)]
        node_sequences[source_index] = [source_index]

        while True:
            current_index = min(unvisited, key=lambda i: distances_from_start[i])

            if current_index == target_index:
                break

            unvisited.remove(current_index)

            for edge in self.edges_adjacent_to(current_index):
                start = current_index
                end = edge.end if edge.start == current_index else edge.start

                if end in unvisited and distances_from_start[end] > distances_from_start[start] + edge.length:
                    distances_from_start[end] = distances_from_start[start] + edge.length
                    node_sequences[end] = node_sequences[start].copy()
                    node_sequences[end].append(end)

        return node_sequences[target_index]