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

'''forcedirected.py



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

import time
import itertools
import math
import json
from littlegeom import Point, Vector, Segment
import intersect2d
import convexhull



class Vertex(object):
	
	def __init__(self, pos):
		self.pos = Point(pos) # 2D position of the point inside the square
		self.intersect = None # 2D position of the point projected onto the square's edge
		self.segment = None   # Index of the edge the point lies on (in 0..3)
		self.relcoord = None	# Coordinate relative to the starting point of the edge
		self.abscoord = None	# Absolute coordinate
		self.speed = 0.0
		self.unprojected = None


		
class ForceDirectedSquareLayout(object):

	GEOMETRIC 	= 0 	# Geometric centroid (mean of (Xi,Yi))
	CUSTOM		= 1		# My custom formula to find the center
	CONVEX_HULL	= 2		# Geometric centroid of the convex hull
	
	p_max_repulsion_dist = 100.0
	p_repulsion = 10000.0
	p_speedup_factor = 100.0
	p_speed_cap = 10.0
	p_drag_coeff = 1.0

	def __init__(self, width, height, points, centroid_mode = GEOMETRIC):
		self.width = width
		self.height = height
		self.points = [Vertex(p) for p in points]
		self.centroid_mode = centroid_mode
		self.centroid = None
		corners = [(0, 0), (self.width, 0), (self.width, self.height), (0, self.height)]
		self.bounds = [Segment(corners[i], corners[(i+1)%len(corners)]) for i in range(len(corners))]
		self.bounds_start_coords = [0, self.width, self.width+self.height, 2*self.width+self.height, 2*(self.width+self.height)]
		self.bounds_vecs = [Vector(1,0), Vector(0, 1), Vector(-1, 0), Vector(0, -1)]
		
		
	def start(self):
		self.calc_centroid()
		self.calc_intersections()
		self.project()
		
		
	def calc_centroid(self):
		geom_centroid = Point([sum(x)/len(self.points) for x in itertools.izip(*(p.pos for p in self.points))])
		if self.centroid_mode == ForceDirectedSquareLayout.CUSTOM:
			wavg = Point()
			sumofw = 0.0
			for pt in self.points:
				dv = pt.pos - geom_centroid
				d = dv.length()
				w = 1./math.pow(d, 3)
				wavg += pt.pos*w
				sumofw += w
			self.centroid = wavg/sumofw
		elif self.centroid_mode == ForceDirectedSquareLayout.CONVEX_HULL:	
			hull = [Point(x) for x in convexhull.convexHull([p.pos for p in self.points])]
			self.centroid = Point([sum(x)/len(hull) for x in itertools.izip(*hull)])
		else:
			self.centroid = geom_centroid
		
		
	def calc_intersections(self):
		assert self.centroid
		proj_dist = 2.0*max(self.width, self.height) # cast a long enough ray so that we're sure it crosses the image's boundaries
		self.proj_points = []
		for pt in self.points:
			v = pt.pos - self.centroid
			dist_from_centroid = v.normalize()
			ray = v * proj_dist
			p2 = self.centroid + ray
			s = Segment(self.centroid, p2)
			for i, seg in enumerate(self.bounds):
				ret = intersect2d.intersect2D_2Segments(s, seg)
				if ret is not None:
					assert pt.intersect is None
					pt.intersect = Point(ret[0])
					pt.segment = i
		assert all(map(lambda p: p.intersect is not None, self.points))
		
		
	def project(self):
		for pt in self.points:
			seg = self.bounds[pt.segment]
			v = pt.intersect - seg[0]
			pt.relcoord = v.length()
			pt.abscoord = self.bounds_start_coords[pt.segment] + pt.relcoord
		
		
	def unproject(self):
		for pt in self.points:
			seg = self.bounds[pt.segment]
			bv = self.bounds_vecs[pt.segment]
			pt.relcoord = pt.abscoord - self.bounds_start_coords[pt.segment]
			pt.unprojected = seg[0] + bv * pt.relcoord
		
	
	def step(self, elapsed):
		for pt in self.points:
			pt.faccum = 0.0
			
		for p1, p2 in itertools.combinations(self.points, 2):
			d = p2.abscoord - p1.abscoord
			if abs(d) < self.p_max_repulsion_dist:
				factor = self.p_repulsion / (d*d) * (1.0 if d > 0.0 else -1.0)
				p1.faccum -= factor
				p2.faccum += factor
		
		for pt in self.points:
			if pt.faccum != 0.0:
				drag = self.p_drag_coeff * pt.speed
				pt.faccum -= drag
				pt.speed += pt.faccum * elapsed * self.p_speedup_factor
				if pt.speed > 0.0: pt.speed = min(pt.speed, self.p_speed_cap)
				else: pt.speed = max(pt.speed, -self.p_speed_cap)
				pt.abscoord += pt.speed * elapsed * self.p_speedup_factor
			else:
				pt.speed = 0.0
	
	def get_energy(self):
		total = 0.0
		for pt in self.points:
			e = 0.5 * pt.speed * pt.speed
			total += e
		return total
		
		
if __name__ == '__main__':
	import pygame
	pygame.init()
	
	input = 'exp_data.json'
	with file(input, 'rb') as f:
		exp_data = json.load(f)
	
	win_size = (800, 600)
	vwin_size = Vector(win_size)
	header_height = 30
	bottom_height = 100
	
	screen = pygame.display.set_mode(win_size)
	black = (0, 0, 0)
	red = (255, 0, 0)
	green = (0, 255, 0)
	blue = (0, 0, 255)
	white = (255, 255, 255)
	selblue = (128, 128, 255)
	yellow = (255, 255, 0)
	
	class Button(object):
		
		ACTIVE = 1
		HIGHLIGHTED = 2
		
		def __init__(self, label, pos, size, data=None):
			self.pos = Point(pos) # make a copy
			self.size = Vector(size) 
			self.rect = pygame.Rect(pos, size)
			font = pygame.font.Font(None, 16)
			self.text = font.render(label, 1, blue)
			self.state = 0
			self.data = data
			
		def draw(self, screen):
			if self.state & Button.ACTIVE:
				pygame.draw.rect(screen, selblue, self.rect, 0)
			color = yellow if self.state & Button.HIGHLIGHTED else blue
			pygame.draw.rect(screen, color, self.rect, 1)
			screen.blit(self.text, self.pos + Vector(5,4))
			
		def hittest(self, point):
			return self.rect.collidepoint(point)
		
		

	def pairwise(iterable):
		"s -> (s0,s1), (s1,s2), (s2, s3), ..."
		# http://docs.python.org/2/library/itertools.html
		a, b = itertools.tee(iterable)
		next(b, None)
		return itertools.izip(a, b)	
	
	class Simulation(object):
		
		def __init__(self):
			self.layout = None
			
		@staticmethod
		def calc_fit_size(obj_size, view_size):
			#width, height = [float(x) for x in image.size]
			objr = float(obj_size[0]) / float(obj_size[1])
			viewr = float(view_size[0]) / float(view_size[1])
			if objr > viewr:
				w = float(view_size[0])
				h = w / objr
			else:
				h = float(view_size[1])
				w = h * objr
			return w,h
			
			
		def is_active(self):
			return self.layout is not None
			
			
		def start(self, image, config):
			print "Starting '%s-%s'" % (image,config)
			data = exp_data[image]['configs'][config]
			if data['name'] == 'hull': mode = ForceDirectedSquareLayout.CONVEX_HULL
			elif data['name'] == 'weighted': mode = ForceDirectedSquareLayout.CUSTOM
			else: mode = ForceDirectedSquareLayout.GEOMETRIC
			self.layout = ForceDirectedSquareLayout(data['size'][0], data['size'][1], data['minutiae'], mode)
			self.layout.start()
			self.transforms = [None,None]
			self.transform_index = None
			self.layout.unproject()
			
		
		def step(self, elapsed):
			self.layout.step(elapsed)
			self.layout.unproject()
			
			
		def set_transform(self, scenesz, pos, size, margins):
			assert self.transform_index is not None
			sz = scenesz + margins
			w, h = Simulation.calc_fit_size(sz, size)
			scale = w / sz[0]
			offset = pos+(size-Vector(w,h))/2.0+(margins/2.0)*scale
			self.transforms[self.transform_index] = offset, scale
			
			
		def draw_top(self, screen, pos, size):
			pygame.draw.rect(screen, blue, pygame.Rect(pos, size), 1)
			self.transform_index = 0

			if self.transforms[0] is None:
				scenesz = Vector(self.layout.width,self.layout.height)
				self.set_transform(scenesz, pos, size, Vector(100,100))
			
			self.draw_proj_rect(screen, red, (0,0), (self.layout.width, self.layout.height), 1)
			self.draw_proj_circle(screen, green, self.layout.centroid, 5, 1)
			for pt in self.layout.points:
				self.draw_proj_line(screen, yellow, pt.pos, pt.unprojected)
				self.draw_proj_line(screen, green, pt.pos, pt.intersect)
				self.draw_proj_circle(screen, red, pt.pos, 5, 1)
				self.draw_proj_circle(screen, yellow, pt.unprojected, 5, 1)
				self.draw_proj_circle(screen, green, pt.intersect, 5, 1)
			
			
		def draw_bottom(self, screen, pos, size):
			pygame.draw.rect(screen, blue, pygame.Rect(pos, size), 1)
			self.transform_index = 1

			if self.transforms[1] is None:
				scenesz = Vector(self.layout.bounds_start_coords[4], 30)
				self.set_transform(scenesz, pos, size, Vector(100,100))
				
			for x1, x2 in pairwise(self.layout.bounds_start_coords):
				self.draw_proj_line(screen, red, Vector(x1,0), Vector(x2,0), 1)
				self.draw_proj_line(screen, red, Vector(x1, -10), Vector(x1, 10), 1)
				self.draw_proj_line(screen, red, Vector(x2, -10), Vector(x2, 10), 1)
				
			for pt in self.layout.points:
				self.draw_proj_circle(screen, green, Vector(pt.abscoord,0), 5, 1)
				
		def draw_proj_circle(self, surface, color, pos, radius, *args, **kwargs):
			pos2 = [int(round(k)) for k in (self.transforms[self.transform_index][0]+(Point(pos)*self.transforms[self.transform_index][1]))]
			radius2 = int(round(radius*self.transforms[self.transform_index][1]))
			if radius2 <= 0: radius2 = 1
			pygame.draw.circle(surface, color, pos2, radius2, *args, **kwargs)
		
		
		def draw_proj_rect(self, surface, color, pos, size, *args, **kwargs):
			pygame.draw.rect(surface, color, pygame.Rect(self.transforms[self.transform_index][0]+Point(pos), Vector(size)*self.transforms[self.transform_index][1]), *args, **kwargs)
			
		def draw_proj_line(self, surface, color, p1, p2, *args, **kwargs):
			P1, P2 = [[int(round(k)) for k in (self.transforms[self.transform_index][0]+(Point(p)*self.transforms[self.transform_index][1]))] for p in (p1,p2)]
			pygame.draw.line(surface, color, P1, P2, *args, **kwargs)
			
			
	buttons = []
	p = Point(5, 5)
	sz = Vector(100,20)
	for image, idata in exp_data.iteritems():
		for config, cdata in idata['configs'].iteritems():
			button = Button('%s-%s' % (image, config), p, sz, data=(image, config))
			p += Vector(sz[0]+5.0, 0.0)
			buttons.append(button)
	buttons.append(Button('quit (ESC)', (win_size[0]-5-sz[0],p[1]), sz, data='quit'))
	running_button = Button('start / stop', (5, header_height+5), sz, data='toggle')
	buttons.append(running_button)
	
	sim = Simulation()
	to_start = ('mark', 'geom')
	filter(lambda b: b.data==to_start, buttons)[0].state |= Button.ACTIVE
	sim.start(*to_start)
	
	def pygame_loop():
		last_time = None
		running = False
		while 1:
			for event in pygame.event.get():
				if event.type == pygame.QUIT:
					return
				elif event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE:
					return
				
				if event.type in (pygame.MOUSEBUTTONDOWN, pygame.MOUSEBUTTONUP, pygame.MOUSEMOTION):

					for btn in buttons:
						btn.state &= ~Button.HIGHLIGHTED
					
					hit = None
					for button in buttons:
						if button.hittest(event.pos):
							hit = button
							hit.state |= Button.HIGHLIGHTED
							break
					
					if hit is not None:
						if event.type == pygame.MOUSEBUTTONDOWN:
							if hit.data == 'quit':
								return
							elif hit.data == 'toggle':
								running = not running
								if running: running_button.state |= Button.ACTIVE
								else: running_button.state &= ~Button.ACTIVE
							else:
								sim.start(*hit.data)
								for btn in buttons[:6]:
									btn.state &= ~Button.ACTIVE
								hit.state |= Button.ACTIVE
								running = False
							
							
			now = time.time()
			if last_time is not None and last_time != now:
				elapsed = now - last_time
			else:
				elapsed = 0.0
				
				
			screen.fill(black)
			
			for button in buttons:
				button.draw(screen)
			
			if sim.is_active():
			
				if running and elapsed > 0.0:
					sim.step(elapsed)
					e_total = sim.layout.get_energy()
					e_per_point = e_total / float(len(sim.layout.points))
					print "Energies: total=%f, per point=%f" % (e_total, e_per_point)
					if e_total == 0.0:
						print "Done!"
						running = False
						running_button.state &= ~Button.ACTIVE
						
				sim.draw_top(screen, Point(0, header_height), Vector(win_size[0], win_size[1]-header_height-bottom_height))
				sim.draw_bottom(screen, Point(0, win_size[1]-bottom_height), Vector(win_size[0], bottom_height))
			
			pygame.display.update()
			
			#energy = fdl.get_energy()
			#print "Total system energy=",energy,"(avg per node=",(energy/len(pts)),")"
			
			last_time = now
	
	pygame_loop()
	