import math
import re

from SoSe24.lec06_graph.disjoint import DisjointValue
from graph import Graph, AdjacencyListGraph, AdjacencyMatrixGraph, Vertex
import heapq


def mst_prim(self, start_name: str = None):
    """ Compute the minimum spanning tree of the graph using Prim's algorithm. """

    distance_map = {}   # maps vertices to their current distance from the spanning tree
    parent_map = {}     # maps vertices to their predecessor in the spanning tree

    Vertex.__lt__ = lambda self, other: distance_map[self] < distance_map[other]

    queue = []

    if start_name is None:
        start_name = self.all_vertices()[0].value;

    # Initialize the maps
    for vertex in self.all_vertices():
        distance_map[vertex] = 0 if vertex.value == start_name else math.inf
        parent_map[vertex] = None
        queue.append(vertex)

    heapq.heapify(queue)    # Convert the list into a heap

    # Process the queue
    cost = 0 # The cost of the minimum spanning tree
    while len(queue) > 0:
        vertex = heapq.heappop(queue)
        cost += distance_map[vertex] # Add the cost of the edge to the minimum spanning tree
        for (dest, w) in self.get_adjacent_vertices_with_weight(vertex.value):
            if dest in queue and distance_map[dest] > w:
                # Update the distance and parent maps
                queue.remove(dest)
                distance_map[dest] = w
                parent_map[dest] = vertex
                queue.append(dest)      # Add the vertex back to the queue
                heapq.heapify(queue)    # Re-heapify the queue

    # Return the distance and predecessor maps
    return parent_map, cost


def mst_kruskal(self, start_name: str = None):
    """ Compute the minimum spanning tree of the graph using Kruskal's algorithm. """

    cost = 0
    result = []
    edges = self.all_edges()

    # Create a disjoint set for each vertex
    vertex_map = {v.value: DisjointValue(v) for v in self.all_vertices()}

    # Sort the edges by weight
    edges.sort(key=lambda edge: edge[2])

    # Process the edges
    for edge in edges:
        start_name, end_name, weight = edge
        # Check if the edge creates a cycle
        if not vertex_map[start_name].same_set(vertex_map[end_name]):
            result.append(edge)
            vertex_map[start_name].union(vertex_map[end_name])
            cost += weight

    return result, cost



AdjacencyListGraph.mst_prim = mst_prim
AdjacencyMatrixGraph.mst_prim = mst_prim

AdjacencyListGraph.mst_kruskal = mst_kruskal
AdjacencyMatrixGraph.mst_kruskal = mst_kruskal


def read_elektro_into_graph(graph: Graph, filename: str):
    pattern = re.compile(r'"([^"]+)";"([^"]+)";(\d+)')
    with (open(filename, "r") as file):
        for line in file:
            m = pattern.match(line)
            if m:
                start_name = m.group(1)
                end_name = m.group(2)
                cost = int(m.group(3))
                graph.insert_vertex(start_name)
                graph.insert_vertex(end_name)
                graph.connect(start_name, end_name, cost)
                graph.connect(end_name, start_name, cost)

if __name__ == "__main__":

    graph = AdjacencyMatrixGraph()
    read_elektro_into_graph(graph, "../../elektro.txt")

    parents, cost = graph.mst_prim()
    print(f"Kosten nach Prim: {cost}")
    for node, parent in parents.items():
        if parent is not None:
            print(f"{node} - {parent}")

    edges, cost = graph.mst_kruskal()
    print(f"Kosten nach Kruskal: {cost}")
    for start_name, end_name, _ in edges:
        print(f"{start_name} - {end_name}")