#
#  Minimum Spanning Tree
#
import random, sys, sarg
from   graphics import GraphWin, Text, Point, Line

dbg = False
colors= ("red","green","blue","orange","yellow","purple")
ncolors = len(colors)

win = GraphWin("Minimum Spanning Tree",400,400)

class Minnode :
    def __init__ (self, pos, label, group) :
        self.pos = pos
        self.label = label
        self.group = group

class Minspan :
    def __init__ (self, pixels, nNodes) :
        letters    = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
        self.drawn = []
        self.nodes = [Minnode(randpos(50,pixels-50), letters[i], i+1)
                                    for i in range(nNodes)]
        self.plinks = []          # potential links between nodes
        self.pixels = pixels
        for m in range(nNodes) :
            for n in range(m+1,nNodes) :
                x = self.nodes[m]
                y = self.nodes[n]
                dist = euclDist2(x.pos,y.pos)
                self.plinks.append((dist,m*n,x,y)) # m*n unique in list
        self.plinks.sort()       # so that sort won't try to compare x and y
        self.links = []

    def linkNodes (self) :
        nLinks = 0
        while self.plinks :
            plink = self.plinks.pop(0)
            dst, n, nodA, nodB = plink
            if nodA.group != nodB.group :
                self.links.append((nodA,nodB))
                masterGroup = nodA.group
                slaveGroup  = nodB.group
                self.display("Will link %s->%s" % (nodA.label, nodB.label))
                self.links[-1] = (nodA,nodB)
                for node in self.nodes :
                    if node.group == slaveGroup :
                        node.group = masterGroup
                self.display("Grouped %s with %s" % (slaveGroup,masterGroup))
                nLinks += 1
        return nLinks

    def display (self,banner) :
        for obj in self.drawn : obj.undraw()  # erase the board
        self.drawn = []                       # done
        if dbg : print("Banner %s" % banner)
        txtBanner = Text(Point(200,20), banner)
        txtBanner.draw(win)
        self.drawn.append(txtBanner)
        for a,b in self.links :
            if a.group != b.group : color = "white"
            else                  : color = colors[a.group%ncolors]
            if dbg : print("line: %s from %s to %s" % (color, a.pos, b.pos))
            p1 = Point(a.pos[0], a.pos[1])
            p2 = Point(b.pos[0], b.pos[1])
            line = Line(p1,p2)
            line.setFill(color)
            line.draw(win)
            self.drawn.append(line)
        for node in self.nodes :
            color = colors[node.group%ncolors]
            txtLabel = Text(Point(node.pos[0],node.pos[1]), node.label)
            txtLabel.setFill(color)
            txtLabel.draw(win)
            self.drawn.append(txtLabel)
            if dbg : print("text: %s %s at %s" % (color, node.label, node.pos))
        if dbg : print()
        win.getMouse()     # wait for the click

def euclDist2(a,b) :
    return (a[0]-b[0])**2 + (a[1]-b[1])**2

def randpos(low,high) :
    from random import randint
    return (randint(low,high), randint(low,high))

def main()  :
    seed    = sarg.Int("seed", 0)
    if seed : random.seed(seed)
    else :
        seed = random.randint(101,999)
        print("Seed is %s" % seed)
        random.seed(seed)
    pixels  = sarg.Int("pixels",400)
    nNodes  = sarg.Int("nodes",   5)
    print("Click the mouse in the window to advance frames")
    ms = Minspan(pixels, nNodes)
    nLinks = ms.linkNodes()

if __name__ == "__main__" : main()