'''
Minimum Spanning Tree generation (SVG) for Prim's algorithm.
Firstly use this code to generate SVG frames.
Then transform to bitmaps and convert to GIF.
'''

# range size
N = 300
margin = 20

def norm(x, y):
    return (x*x y*y)**.5

def saveToSVG(nFrames, points, firmed, trying):
    f = open('demo_' '0'*(3-len(str(nFrames))) str(nFrames) '.svg', 'w')
    f.write("<svg xmlns=\"http://www.w3.org/2000/svg\" version=\"1.1\">\n")
    for p in points:
        f.write("<circle cx=\""  str(p[0] margin)  "\" cy=\""  str(N-p[1] margin)  "\" r=\"5\" fill=\"white\" stroke=\"black\"/>\n")
    for L in firmed:
        f.write("<line x1=\""  str(L[0][0] margin)  "\" y1=\""  str(N-L[0][1] margin)  "\" x2=\""   str(L[1][0] margin)   "\" y2=\""   str(N-L[1][1] margin)   "\" stroke=\"red\" stroke-width=\"5\"/>\n")
    for L in trying:
        f.write("<line x1=\""  str(L[0][0] margin)  "\" y1=\""  str(N-L[0][1] margin)  "\" x2=\""   str(L[1][0] margin)   "\" y2=\""   str(N-L[1][1] margin)   "\" stroke=\"blue\" stroke-width=\"5\"/>\n")
    f.write("</svg>\n")
    f.close()

def generatePoints(n):
    import random as r
    r.seed(100)
    
    res = []
    for i in range(n):
        pt = [r.randint(0,N) for _ in [0, 1]]
        if [pt] not in res:
            res  = [pt]
    return res

def prim(n, points):
    n = len(points)
    mst = []
    S = set()
    import heapq
    heap = []
    nframe = 0
    while len(mst)<n-1:
        if len(S)==0:
            topV = 0
        else:
            while heap[0][2] in S:
                heapq.heappop(heap)
            topV = heap[0][2]
            saveToSVG(nframe, points, mst, [[points[heap[0][1]], points[heap[0][2]]]])
            nframe =1
            mst.append([points[heap[0][1]], points[topV]])
            heapq.heappop(heap)
            saveToSVG(nframe, points, mst, [])
            nframe =1
        S.add(topV)
        for i in range(n):
            if i not in S:
                L = norm(points[i][0]-points[topV][0], points[i][1]-points[topV][1])
                heapq.heappush(heap, [L, topV, i])
    return mst

# test 30 points temporarily
n = 30
pts = generatePoints(n)
prim(n, pts)