graph

from typing import TypeVar, Generic, Callable
from typing import Tuple
import pickle
from dataclasses import dataclass

T = TypeVar("T")

@dataclass

class Graph(Generic[T]):
    def __init__(self, is_directed: bool = False, ) -> None: #задаем данные матрицы
        self.length: int = 0
        self.matrix: list[list[int]] = []
        self.vertex_names: dict[T, int] = {}
        self.is_directed: bool = is_directed

    def add_vertex(self,  name: T) -> int: #добавляем вершину
        index = self.length
        self.matrix.append([0] * self.length)
        for row in self.matrix:
            row.append(0)
        self.length += 1
        self.vertex_names[name] = index
        return index

    def add_edge(self, v1: int, v2: int, weight: int) -> None: #добавляем грань
        if v1 >= self.length or v2 >= self.length:
            raise Exception("out of range")
        self.matrix[v1][v2] = weight
        if not self.is_directed:
            self.matrix[v2][v1] = weight

    def add_without_weight(self, vertex1: T, vertex2: T) -> None: #добавляем грань без веса
        self.add_edge(vertex1, vertex2, 0)

    def for_each_vertex(self, fun: Callable[[T], None]) -> None: #применить функцию для каждой вершины
        for vert, _ in self.matrix:
            fun(vert)

    def amount_vertexes(self) -> int: #колчиество вершин
        return len(self.vertex_names)

    def for_each_adj_edge(self, vertex: T, fun: Callable[[int], None]): #применить функцию для каждой грани
        if vertex in self.vertex_names:
            edges: list[int] = self.matrix[self.vertex_names[vertex]]
            for number in edges:
                fun(number)

    def print_matrix(self) -> None: #вывести матрицу смежности
        print("    ", end="")
        for index in range(self.length):
            print(f"{self.get_vertex_by_index(index):<8}", end="")
        print()

        for i, row in enumerate(self.matrix):
            print(f"{self.get_vertex_by_index(i):<4}", end="")
            for val in row:
                print(f"{val:<8}", end="")
            print()

    def get_vertex_by_index(self, index: int) -> T: #вытащить вершину по ее номеру
        for vertex, vert_index in self.vertex_names.items():
            if vert_index == index:
                return vertex
        raise ValueError("Vertex not found")

    def prim(self) -> list[Tuple[T, T, int]]: #поиск минимального остовного дерева
        if self.length == 0:
            return []

        visited = set()
        edges = []
        start_vertex = next(iter(self.vertex_names))
        visited.add(start_vertex)

        while len(visited) < self.length:
            min_edge = None
            min_weight = float('inf')

            for visited_vertex in visited:
                visited_vertex_index = self.vertex_names[visited_vertex]

                for neighbor_index, weight in enumerate(self.matrix[visited_vertex_index]):
                    if weight != 0 and self.get_vertex_by_index(neighbor_index) not in visited:
                        if weight < min_weight:
                            min_edge = (visited_vertex, self.get_vertex_by_index(neighbor_index), weight)
                            min_weight = weight

            if min_edge:
                edges.append(min_edge)
                visited.add(min_edge[1])

        return edges

    def save_to_file(self, filename: str) -> None:
        data = {
            'length': self.length,
            'matrix': self.matrix,
            'vertex_names': self.vertex_names,
            'is_directed': self.is_directed
        }

        with open(filename, 'wb') as file:
            pickle.dump(data, file)

    def load_from_file(self, filename: str) -> None:
        with open(filename, 'rb') as file:
            data = pickle.load(file)

        self.length = data['length']
        self.matrix = data['matrix']
        self.vertex_names = data['vertex_names']
        self.is_directed = data['is_directed']