from dataclasses import dataclass
from typing import List


@dataclass
class Edge:
    weight: int
    source_id: int
    target_id: int


@dataclass
class Vertex:
    id: int
    edges: List[Edge]


def load_matrix(filename):
    def line_to_row(line):
        r = []
        for field in line.strip().split(","):
            if field == "-":
                r.append(0)
            else:
                r.append(int(field))
        return r

    with open(filename, "r") as f:
        m = [line_to_row(l) for l in f]
    return m


def matrix_to_graph(matrix):
    n_nodes = len(matrix)
    graph = []
    for row_index in range(n_nodes):

        edges = []
        for col_index in range(n_nodes):
            weight = matrix[row_index][col_index]
            if weight != 0:
                edges.append(Edge(weight, row_index, col_index))
        graph.append(Vertex(row_index, edges))
    return graph


def total_weight(graph):
    total_weight = 0
    for vertex in graph:
        for edge in vertex.edges:
            total_weight += edge.weight
    total_weight //= 2
    return total_weight


def prims_algo(graph):
    """ Takes a graph and returns the edges that for a minimum spanning tree
    for that graph. """
    visited_id = set([0])
    min_edges = []
    while len(visited_id) != len(graph):
        min_edge = None
        for vertex_id in visited_id:
            vertex = graph[vertex_id]
            for edge in vertex.edges:
                if edge.target_id in visited_id:
                    pass
                elif min_edge is None:
                    min_edge = edge
                elif edge.weight < min_edge.weight:
                    min_edge = edge
        visited_id.add(min_edge.target_id)
        min_edges.append(min_edge)
    return min_edges


def euler_107():
    m = load_matrix("../txt/e107.txt")
    g = matrix_to_graph(m)
    weight = total_weight(g)
    new_weight = sum([e.weight for e in prims_algo(g)])
    return weight - new_weight


if __name__ == "__main__":
    solution = euler_107()
    print("e107.py: " + str(solution))
    assert(solution == 259679)