#!/usr/bin/python
# -*- coding: utf-8 -*-

'''test_forcedirected2d.py

Preliminary test of generic 2D force-directed layout.

Depends on:
	littlegeom.py
	
Author: 
	Thibault Genessay
'''

from littlegeom import Point, Vector
import time
import itertools
import math

class Vertex(object):

	def __init__(self, pos):
		self.pos = pos
		self.speed = Vector()
		self.faccum = None
		self.anchored = False
		
class FDL(object):
	
	def __init__(self, points, links, factor=1.0):
		#self.points = points
		self.links = links
		self.vertices = []
		for p in points:
			self.vertices.append(Vertex(p))
		self.factor = factor
			
	def step(self, elapsed):
		# Reset the force accumulator for all vertices
		for v in self.vertices:
			v.faccum = Vector()

		# Calculate link spring forces (each link is a spring of length 'spring_length'
		# that applies a force F = k.V where k has been determined by trial-and-error,
		# and V is the normalized vector between the two vertices. The force will push
		# the vertices away if the length is > spring_length and pull them together if
		# the length is < spring_length.
		spring_length = 150.0
		for link in self.links:
			v1, v2 = [self.vertices[x] for x in link]
			v = v2.pos - v1.pos
			lv = v.normalize()
			target_lv = spring_length
			diff = lv - target_lv
			fmag = diff*10.0
			f = v * fmag
			v1.faccum += f
			v2.faccum -= f
			
		# Calculate mutual repulsion forces. This is independent of the links and acts
		# on all possible pairs of vertices. The force is proportional to the inverse
		# of the distance between 2 nodes - a bit like electronic atomic repulsion
		# forces.
		for v1, v2 in itertools.combinations(self.vertices, 2):
			d = v2.pos - v1.pos
			dl = d.normalize()
			if dl < 50.0:
				fmag = 10000.0/dl
				f = d * fmag
				v1.faccum -= f
				v2.faccum += f
	
		# Apply the forces resulting from the calculation to the vertices.
		# The speed is capped and damped to faster the convergence and
		# prevent oscillation.
		for v in self.vertices: 
			if not v.anchored:
				v.speed += v.faccum * elapsed
				speed = v.speed.length()
				if speed > 0.0:
					v.speed /= speed
					# Max speed capping
					speed = min(speed, 20.0)
					v.speed = v.speed * speed
					v.pos += v.speed * elapsed * self.factor
					# Speed damping
					v.speed *= 0.9
	
	
	def get_energy(self):
		total = 0.0
		for v in self.vertices:
			if not v.anchored:
				e = 0.5 * math.pow(v.speed.length(), 2.0)
				total += e
		return total
		
		
if __name__ == '__main__':
	import pygame
	pygame.init()
	
	win_size = (800, 600)
	screen = pygame.display.set_mode(win_size)
	black = (0, 0, 0)
	blue = (0, 0, 255)
	
	import random
	pts = []
	for i in range(60):
		x = random.uniform(0, 800)
		y = random.uniform(0, 600)
		pt = Point(x,y)
		pts.append(pt)
	
	links = []
	
	# Random links
	# for i in range(300):
		# p1 = random.randint(0, len(pts)-1)
		# p2 = random.randint(0, len(pts)-1)
		# while p2 == p1:
			# p2 = random.randint(0, len(pts)-1)
		# links.append((p1,p2))

	# Star-shaped graph (links from node 0 to node i, for i in 1..n)
	for i in range(1, len(pts)):
		links.append((0,i))
		
	fdl = FDL(pts, links, 3.0)
	fdl.vertices[0].pos = Point([x/2.0 for x in win_size])
	fdl.vertices[0].anchored = True

	shouldContinue = True
	last_time = None
	while shouldContinue:
		for event in pygame.event.get():
			if event.type == pygame.QUIT:
				shouldContinue = False
				break
				
		if not shouldContinue:
			pygame.quit()
			break

		now = time.time()
		if last_time is not None and last_time != now:
			elapsed = now - last_time
			fdl.step(elapsed)
			
		screen.fill(black)
		
		for pt in fdl.vertices:
			pygame.draw.circle(screen, blue, [int(x) for x in pt.pos], 5, 0)
		for link in fdl.links:
			p1, p2 = [fdl.vertices[x] for x in link]
			pygame.draw.line(screen, blue, [int(x) for x in p1.pos], [int(x) for x in p2.pos])
			
		pygame.display.update()
		
		energy = fdl.get_energy()
		print "Total system energy=",energy,"(avg per node=",(energy/len(pts)),")"
		
		last_time = now
	