#!/usr/bin/env python3

"""
Simple implementation of Prims's algorithm with step-by-step
output of intermediate results in GraphViz.

Input cam be either a well-behaved GraphViz graph, or a simple list of
weighted edges. Warnings of the form "Cannot handle X" are usually
benevolent and can probably be ignored.

Note to Coders: For optimal results, the output of all GraphViz tools
should order nodes in ascending order, have all edges be from smaller
to larger node, and have edges be in ascending order. This assures (?)
uniform layout of equivalent graphs.
"""

import sys
import re
import random
import getopt



Infty = 1000000000

class Node(object):
    def __init__(self, id):
        self.id = id
        self.connected  = False

    def nodestr(self):
        color = ",style=filled,color=\"#D0D0D0\""
        if self.connected:
            color = ",style=filled,color=\"#FFA1A3\""
        node = "   %s [label = <%s>%s]\n"%(self.id,self.id, color)
        return node

    def __str__(self):
        return self.nodestr()

    def __lt__(self, other):
        return self.id < other.id

    def __le__(self, other):
        return self.id <= other.id

    def __gt__(self, other):
        return self.id > other.id

    def __ge__(self, other):
        return self.id >= other.id


class Graph(object):
    def __init__(self, nodes, edges = []):
        self.nodes = {}
        self.nodelist = []
        self.edges = {}
        self.used_edges = set()
        for i in nodes:
            self.addNode(i)
        for e in edges:
            self.addEdge(e)

    def addNode(self, id):
        try:
            node = self.nodes[id]
        except KeyError:
            node = Node(id)
            self.nodes[id] = node
            self.nodelist.append(id)
        return node

    def getNodeNames(self):
        return list(self.nodes.keys())

    def node(self,id):
        return self.nodes[id]

    def addEdge(self, edge):
        (node1, node2),weight = edge
        n1 = self.addNode(node1)
        n2 = self.addNode(node2)
        self.edges[(n1, n2)] = weight
        self.edges[(n2, n1)] = weight

    def getCandEdges(self):
        res = []
        for v1, v2 in self.edges:
            if v1.connected and not v2.connected:
                res.append((self.edges[(v1, v2)], v1, v2))
        return res

    def useEdge(self, v1, v2):
        if v2.id < v1.id:
            v1,v2 = v2,v1
        self.used_edges.add((v1, v2))
        v1.connected = True
        v2.connected = True


    def __str__(self):
        nodes = sorted([str(i) for i in self.nodes.values()])
        edges = []
        for v1, v2 in self.edges:
            if v1.id < v2.id:
                if (v1,v2) in self.used_edges:
                    colour = ",color=\"red\""
                else:
                    colour = ",color=\"grey\""
                edge = v1.id+" -> "+v2.id+" [label=\"%d\"%s]"%(self.edges[(v1,v2)],colour)

                edges.append(edge)

        edges.sort()
        return """
digraph mygraph {
   rankdir = LR
   #bgcolor=transparent
   node [fontsize=18]
   edge [fontsize=14, dir=none]

"""\
    + "\n".join(nodes)+"\n"\
    + "\n".join(edges)+"""
}
"""


def prim(graph, start):
    writeGraph(graph)
    total_weight = 0
    graph.node(start).connected = True

    while True:
        cands = graph.getCandEdges()
        if not cands:
            break
        weight, v1, v2 = min(cands)
        total_weight+=weight
        if v1.connected:
            new = v2.id
        else:
            new = v1,id
        print("Adding %s with %d (total weight %d)"%(new,weight,total_weight))
        graph.useEdge(v1, v2)
        writeGraph(graph)

counter = 0

def newfile():
    global counter
    name = "prim%03d.gv"%(counter,)
    counter = counter+1
    print("Writing to file", name)
    return open(name, "w");

def writeGraph(graph):
    fp = newfile()
    print(graph, file=fp)
    fp.close()


gvvertexRe = re.compile("\s*(\w+)\s*->\s*(\w+)\s*\[label=\"\s*(\d*)\s*\"\]")
gvedgeRe   = re.compile("\s*(\w+)\s*\[label")


def parseGraph(name):
    fp = open(name, "r")
    res = []
    nodes = []
    for line in fp:
        mo1 = gvvertexRe.match(line)
        mo2 = gvedgeRe.match(line)
        if mo1:
            print("Match", line)
            f,t,w = mo1.groups()
            res.append(((f.strip(), t.strip()), int(w)))
        elif mo2:
            nodes.append(mo2.groups()[0])
        else:
            try:
                e,w=line.split(":")
                f,t=e.split("->")
                res.append(((f.strip(), t.strip()), int(w)))
            except ValueError:
                print("Cannot handle ", line)
    fp.close()
    return nodes, res


if __name__ == '__main__':
    opts, args = getopt.gnu_getopt(sys.argv[1:],
                                   "h",
                                   ["help"])
    print(args)
    start = "a"
    if args:
        nodes, edges = parseGraph(args[0])
        graph = Graph(nodes, edges)
        if len(args)>=2:
            start = args[1]
        else:
            start= graph.getNodeNames()[0]

    else:
        graph = Graph([], [
            (("a", "b"), 17),
            (("a", "c"), 8),
            (("a", "f"), 6),
            (("a", "g"), 12),
            (("b", "c"), 4),
            (("b", "d"), 2),
            (("b", "e"), 15),
            (("c", "e"),  5),
            (("c", "g"),  9),
            (("d", "f"), 10),
            (("e", "f"), 14),
            (("e", "g"), 11),
            (("f", "g"), 16)])

    prim(graph, start)
    # writeGraph(graph)