# David Baird
# 2004/05/11
# My implementation of neato, part of graphviz.
#
# Basic algorithm:
# 1. Place an ideal spring between each pair of nodes
# 2. Set the spring unstretched distance equal to the shortest distance
#    path between node pairs in the graph
# 3. Run an iterative solver to model the springs, but don't track the
#    velocity state
# 4. Tada!
#
# Caveats:
# * Don't initialize node locations to be all on the same line, otherwise
#   they will likely always remain on this straight line
#
# This file is a little messy because I was in the middle of converting
# from bland python lists to numarray.
#

import numarray
from numarray import array as Array
import Tkinter
import time
from math import sqrt, fabs

program = "Experiments in Dynamics with Python and a Tk Canvas"

class GraphView(Tkinter.Frame):
    def __init__(self, nodes, edges, master=None):
        Tkinter.Frame.__init__(self, master)
        self.grid()
        self.radius = 5
        self.__createWidgets()
        self.__createNodes(nodes=nodes, edges=edges, radius=self.radius)
    def __createWidgets(self):
        self.canvas = Tkinter.Canvas(self, width=400, height=400)
        self.canvas.grid()
    def __createNodes(self, nodes, edges, radius):
        self.node_ids = []
        self.edge_ids = []
        for node in nodes:
            x, y = node
            id = self.canvas.create_oval(x-radius, y-radius, x+radius,
                                         y+radius)
            self.node_ids.append(id)
        for edge in edges:
            a, b = edge
            x1, y1 = nodes[a]
            x2, y2 = nodes[b]
            id = self.canvas.create_line(x1, y1, x2, y2)
            self.edge_ids.append(id)
    def updateNodes(self, nodes, edges):
        #for kaka in nodes, self.node_ids:
        r = self.radius
        for i in range(len(nodes)):
            x, y = nodes[i]
            id = self.node_ids[i]
            #self.canvas.move(id, x, y)
            self.canvas.coords(id, x-r, y-r, x+r, y+r)
        for i in range(len(edges)):
            id = self.edge_ids[i]
            a, b = edges[i]
            x1, y1 = nodes[a]
            x2, y2 = nodes[b]
            self.canvas.coords(id, x1, y1, x2, y2)

class SpringModel:
    def __init__(self, nodes, springs):
        self.node_count = len(nodes)
        node_range = self.node_range = range(self.node_count)
        self.node_pairs = \
            [(a, b) for a in node_range for b in node_range if b > a]
        self.velocity_x = [0.] * self.node_count
        self.velocity_y = [0.] * self.node_count
    def __forces(self, nodes, springs, k):
        '''Todo: reimplement this in numarray!'''
        force = []
        for i in self.node_range:
            force.append([0., 0.])
        for pair in self.node_pairs:
            a, b = pair
            disp_ab_x = nodes[a][0] - nodes[b][0]
            disp_ab_y = nodes[a][1] - nodes[b][1]
            dist = sqrt(disp_ab_x**2 + disp_ab_y**2)
            eq_dist = springs[pair]
            stretch = dist - eq_dist
            #print pair, dist, eq_dist
            f = -k * stretch
            f_a_x = f * disp_ab_x / dist
            f_a_y = f * disp_ab_y / dist
            force[a][0] += f_a_x
            force[a][1] += f_a_y
            force[b][0] -= f_a_x
            force[b][1] -= f_a_y
        return force
    def __move(self, forces, nodes, mass, dt):
        for i in self.node_range:
            f_x, f_y = forces[i]
            self.velocity_x[i] += f_x * dt / mass
            self.velocity_y[i] += f_y * dt / mass
            nodes[i][0] += self.velocity_x[i] * dt
            nodes[i][1] += self.velocity_y[i] * dt
    def __move2(self, forces, nodes, mass, dt):
        sign = lambda x: x/fabs(x)
        f_friction = lambda f, f2: fabs(f) > fabs(f2) \
            and (-sign(f) * fabs(f2)) \
            or (-f)
        for i in self.node_range:
            f_x, f_y = forces[i]
            nodes[i][0] += f_x * dt / mass
            nodes[i][1] += f_y * dt / mass
    def move(self, nodes, springs):
        #self.__move(self.__forces(nodes, springs, 1), nodes, 1, .01)
        self.__move2(self.__forces(nodes, springs, 1), nodes, 1, .1)


# Initialize a list of nodes and their physical locations:
nodes = map(lambda x: Array(x, type=numarray.numerictypes.Float),
            [[100,100], [150,150], [120,130], [200, 110]])
# Describe the edges of the undirected graph:
edges = [(0, 1), (1, 2), (1, 3), (2, 3)]
# Todo: calculate spring values based on edges:
springs = {
    (0, 1): 100.,
    (0, 2): 200.,
    (0, 3): 200.,
    (1, 2): 100.,
    (1, 3): 100.,
    (2, 3): 100.,
    }

widget = GraphView(nodes, edges)
widget.master.title(program)
widget.update()

solver = SpringModel(nodes, springs)

while 1:
    solver.move(nodes, springs)
    time.sleep(.1)
    widget.updateNodes(nodes, edges)
    widget.update()