# GVB - a GTK+/GNOME vibrations simulator
#
# Copyright (C) 2008 Pietro Battiston
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA



from scipy import shape
from gtk import STATE_NORMAL, ListStore
from gtk.gdk import Pixmap
from time import time
from cairo import ImageSurface, FORMAT_ARGB32, Context

def one_basic(cr, size, pos, speeds = None):
	spaced=1./(len(pos)+1)

	cr.scale(size[0], size[1]/4)
	cr.translate(0, 2)
	cr.set_line_width(.01)
	cr.set_source_rgb(1,0,0)

	cr.move_to(0,0)

	xposition=0
	for ypos in pos:
		xposition=xposition+spaced
#		print xposition, point
		cr.line_to(xposition, -ypos)
	cr.line_to(1,0)
	cr.stroke()


def with_speeds(cr, size, pos, speeds):
	n = len(pos)
	spaced=1./(n+1)
	coeff = .1*n
	cr.scale(size[0], size[1]/4)
	cr.translate(0, 2)
	cr.set_line_width(.1/n)
	cr.set_source_rgb(1,0,0)


	cr.move_to(0,0)

	xposition=spaced
	for idx in range(len(pos)):
		cr.move_to(xposition,-pos[idx])
		cr.rel_line_to(0,-speeds[idx]*coeff)
		cr.stroke()
		xposition=xposition+spaced
#		print xposition, point


	cr.set_line_width(.02)
	cr.set_source_rgb(0,0,1)
	cr.move_to(0,0)

	xposition=0
	for ypos in pos:
		xposition=xposition+spaced
#		print xposition, point
		cr.line_to(xposition, -ypos)
	cr.line_to(1,0)
	cr.stroke()


def temperature(cr, size, pos, speeds = None):
	spaced=1./(len(pos))

	cr.scale(size[0], size[1]/4)
	cr.translate(0, 2)
	cr.set_line_width(.5)
	cr.set_source_rgb(1,0,0)

	cr.move_to(0,0)

	xposition=0
	for ypos in pos:
		cr.move_to(xposition,0)
		xposition=xposition+spaced
#		print xposition, point
		if ypos >= 1:
			cr.set_source_rgb(1,1,1)
		elif ypos > .5:
			cr.set_source_rgb((1-ypos)/2,(1-ypos)/2,0)
		elif ypos > 0:
			cr.set_source_rgb(ypos/2, 1-1.5*ypos, 0)
		elif ypos > -.5:
			cr.set_source_rgb(0,1+2*ypos,-ypos)
		else:
			cr.set_source_rgb(0,0,1+ypos)
		cr.line_to(xposition, 0)
		cr.stroke()

#	print pos[64]



def _two_colours(cr, pos, size, speeds = None):
#	print "shape", shape(pos)
#	print pos
	hspaced=1./shape(pos)[0]
	vspaced=1./shape(pos)[1]

	cr.scale(size[0], size[1])
	yposition=0
	for ypos in pos:
		xposition=0
#		print ypos
		for xpos in ypos:
			if xpos > 0:
				cr.set_source_rgb(1-xpos,1,1)
			else:
				cr.set_source_rgb(1,1+xpos,1)
			cr.rectangle(xposition, yposition, hspaced, vspaced)
			cr.fill()
			xposition=xposition+hspaced
		yposition=yposition+vspaced


def colours(cr, size, pos, speeds = None):
#	print "shape", shape(pos)
#	print pos
	hspaced=1./shape(pos)[0]
	vspaced=1./shape(pos)[1]

	cr.scale(size[0], size[1])
	cr.translate(0, vspaced/2)
	cr.set_line_width(vspaced)
	yposition=0
	for ypos in pos:
		xposition=0
		cr.move_to(0,yposition)
#		print ypos
		for xpos in ypos:
			cr.move_to(xposition, yposition)
			xposition=xposition+hspaced
#			if xpos > 0:
#				cr.set_source_rgb(xpos,(2-xpos)/3,0)

			if xpos >= 1:
				cr.set_source_rgb(1,1,1)
			elif xpos > .5:
				cr.set_source_rgb((1-xpos)/2,(1-xpos)/2,0)
			elif xpos > 0:
				cr.set_source_rgb(xpos/2, 1-1.5*xpos, 0)
			elif xpos > -.5:
				cr.set_source_rgb(0,1+2*xpos,-xpos)
			else:
				cr.set_source_rgb(0,0,1+xpos)

			cr.line_to(xposition, yposition)
			cr.stroke()
		yposition=yposition+vspaced



def trid(cr, size, pos, speeds = None):
#	print "shape", shape(pos)
#	print pos
	m,n=shape(pos)

	vspaced=1./(3*n)

	yposition=.3

	cr.set_line_join(2)

	pos = pos * .5

	tr = ([[yposition for i in range(n+2)]]
		+[[yposition+(j+1)*vspaced]+[yposition + (j+1)*vspaced-pos[j][i]/2 for i in range(n)]+[yposition+(j+1)*vspaced] for j in range(m)]
		+[[yposition+(m+1)*vspaced for i in range(n+2)]])

	cr.scale(size[0], size[1])
#	cr.translate(0, vspaced/2)
	cr.set_line_width(.001)

	top_margin = .1

	for row_idx in range(m+1):
		margin = top_margin*(1-(float(row_idx)/(m)))
		marg_under = top_margin*(1-(float(row_idx+1)/(m)))
		hspaced=(1-2*margin)/(n+3)
		xposition = margin + hspaced
		hsp_under=(1-2*marg_under)/(n+3)
		xpos_under = marg_under + hsp_under
#		print ypos
		for col_idx in range(n+1):
			if tr[row_idx][col_idx+1] > tr[row_idx+1][col_idx+1]:
				cr.set_source_rgba(1,0,0,.8)
			else:
				cr.set_source_rgba(0,1,0,.8)
			cr.move_to(xposition, tr[row_idx][col_idx])
			cr.line_to(xposition + hspaced, tr[row_idx][col_idx+1])
			cr.line_to(xpos_under + hsp_under, tr[row_idx+1][col_idx+1])
			cr.close_path()
			cr.fill()

			if tr[row_idx][col_idx] > tr[row_idx+1][col_idx]:
				cr.set_source_rgba(1,0,0,.8)
			else:
				cr.set_source_rgba(0,1,0,.8)
			cr.move_to(xposition, tr[row_idx][col_idx])
			cr.line_to(xpos_under + hsp_under, tr[row_idx+1][col_idx+1])
			cr.line_to(xpos_under, tr[row_idx+1][col_idx])
			cr.close_path()
			cr.fill()
#			if xpos > 0:
#				cr.set_source_rgb(xpos,(2-xpos)/3,0)

			cr.set_source_rgb(0,0,0)

			cr.move_to(xposition, tr[row_idx][col_idx])
			cr.line_to(xposition + hspaced, tr[row_idx][col_idx+1])
			cr.line_to(xpos_under + hsp_under, tr[row_idx+1][col_idx+1])
			cr.line_to(xpos_under, tr[row_idx+1][col_idx])
			cr.close_path()
			cr.stroke()
			xposition=xposition+hspaced
			xpos_under=xpos_under+hsp_under
#	print vspaced, hspaced

def trid_speeds(cr, size, pos, speeds = None):
#	print "shape", shape(pos)
#	print pos
	m,n=shape(pos)

	vspaced=1./(3*n)

	coeff = .05*max(m,n)

	yposition=.3

	cr.set_line_join(2)

	pos = pos * .5

	tr = ([[yposition for i in range(n+2)]]
		+[[yposition+(j+1)*vspaced]+[yposition + (j+1)*vspaced-pos[j][i]/2 for i in range(n)]+[yposition+(j+1)*vspaced] for j in range(m)]
		+[[yposition+(m+1)*vspaced for i in range(n+2)]])

	cr.scale(size[0], size[1])
#	cr.translate(0, vspaced/2)
	cr.set_line_width(.001)

	top_margin = .1

	for row_idx in range(m+1):
		margin = top_margin*(1-(float(row_idx)/(m)))
		marg_under = top_margin*(1-(float(row_idx+1)/(m)))
		hspaced=(1-2*margin)/(n+3)
		xposition = margin + hspaced
		hsp_under=(1-2*marg_under)/(n+3)
		xpos_under = marg_under + hsp_under
#		print ypos
		for col_idx in range(n+1):
			cr.set_line_width(.001)
			if tr[row_idx][col_idx+1] > tr[row_idx+1][col_idx+1]:
				cr.set_source_rgba(0,0,1,.8)
			else:
				cr.set_source_rgba(0,1,0,.8)
			cr.move_to(xposition, tr[row_idx][col_idx])
			cr.line_to(xposition + hspaced, tr[row_idx][col_idx+1])
			cr.line_to(xpos_under + hsp_under, tr[row_idx+1][col_idx+1])
			cr.close_path()
			cr.fill()

			if tr[row_idx][col_idx] > tr[row_idx+1][col_idx]:
				cr.set_source_rgba(0,0,1,.8)
			else:
				cr.set_source_rgba(0,1,0,.8)
			cr.move_to(xposition, tr[row_idx][col_idx])
			cr.line_to(xpos_under + hsp_under, tr[row_idx+1][col_idx+1])
			cr.line_to(xpos_under, tr[row_idx+1][col_idx])
			cr.close_path()
			cr.fill()
#			if xpos > 0:
#				cr.set_source_rgb(xpos,(2-xpos)/3,0)

			cr.set_source_rgb(0,0,0)

			cr.move_to(xposition, tr[row_idx][col_idx])
			cr.line_to(xposition + hspaced, tr[row_idx][col_idx+1])
			cr.line_to(xpos_under + hsp_under, tr[row_idx+1][col_idx+1])
			cr.line_to(xpos_under, tr[row_idx+1][col_idx])
			cr.close_path()
			cr.stroke()

			if row_idx != 0 and col_idx != 0:
				cr.set_line_width(.1/n)

				cr.set_source_rgba(1,0,0,.5)
				cr.move_to(xposition, tr[row_idx][col_idx])
				cr.rel_line_to(0,-speeds[row_idx-1][col_idx-1]*coeff)
				cr.stroke()

			xposition=xposition+hspaced
			xpos_under=xpos_under+hsp_under
#	print vspaced, hspaced


drawers_1d = {
				'wave' :	one_basic,
				'temp' :	temperature,
				'speeds':	with_speeds
			}

drawers_2d = {
				'geo' :		colours,
				'3D' :		trid,
				'speeds':	trid_speeds
			}


drawers_dict={ 1: drawers_1d, 2: drawers_2d }
drawers = {1: ['wave', 'temp', 'speeds'], 2: ['3D', 'geo', 'speeds']}

class Drawer():
	def __init__(self, drawing, combo, dr_type = None):
		self.drawing = drawing
		if drawing.window:
			self.pixmap = Pixmap(drawing.window, drawing.window.get_size()[0], drawing.window.get_size()[1], -1)
		else:
			self.pixmap = None
		self.combo = combo
		self.active = True
		self.points = None
		self.dim = None
		self.ready = False
		self.drawn = 0
		self.dr_type = dr_type
		self.combo.connect('changed', self.change_drawer)
		if hasattr(self.drawing, 'signals_ids'):
			# This means that the drawing area was already drawn by another
			# drawer. This one is a possessive guy and want her to forget him.
			for signal_id in self.drawing.signals_ids:
				self.drawing.disconnect(signal_id)
		
		# Throw away all his ids! You're mine now!
		self.drawing.signals_ids = []
		self.drawing.signals_ids.append(self.drawing.connect('configure-event', self.draw))
		self.drawing.signals_ids.append(self.drawing.connect('expose-event', self.redraw))
		self.size = None

	def draw(self, *args):
#		print "OK, drawing", self, self.__class__, self.drawer, self.dim
		if self.dim != len(self.points.shape):
			raise
#		else:
#			print "points", self.points.shape, self.points
#		print args
		if not self.drawing.window:
			return

#		if not self.ready:
#			self.configure()


		self.pixmap = Pixmap(self.drawing.window, self.drawing.window.get_size()[0], self.drawing.window.get_size()[1], -1)
		self.pixmap_size = self.drawing.window.get_size()[0:1]
		self.cr = self.pixmap.cairo_create()
		self.cr.set_source_rgb(1,1,1)
		self.cr.paint()

		before=time()
		self.drawer(self.cr, self.drawing.window.get_size(), self.points.pos, self.points.speeds)

		self.drawing.window.clear()
		self.drawing.window.draw_drawable(self.drawing.get_style().fg_gc[STATE_NORMAL], self.pixmap, 0, 0, 0, 0, -1, -1)
		self.drawn=time()
		return self.drawn-before

	def dump_draw(self):

		surface = ImageSurface(FORMAT_ARGB32, self.size[0], self.size[1])

		self.cr = Context(surface)

		before=time()
		self.drawer(self.cr, self.size, self.points.pos, self.points.speeds)

		surface.write_to_png('%s/%08d.png' % (self.dumpdir, self.dumpindex))
		self.dumpindex += 1
		self.drawn=time()
		return self.drawn-before

	def redraw(self, widget, event):
		x, y, width, height = event.area
		if self.pixmap:	#If "expose" wasn't called before than "configure"
			if self.pixmap_size == self.drawing.window.get_size()[0:1]:	 # If we're not redrawing anyway:
				self.drawing.window.draw_drawable(self.drawing.get_style().fg_gc[STATE_NORMAL], self.pixmap, x, y, x, y, width, height)
			else:
				self.draw()

	def change_drawer(self, combo):
		self.dr_type = self.combo.get_active_text()
#		print "change self.dr_type", self.dr_type, self
		if self.dr_type in drawers[self.dim]:
			self.drawer=drawers_dict[self.dim][self.dr_type]
			self.draw()

	def configure(self, points=None, dim=None):
		if not self.active:
			return

#		print combo
		if points:
			self.points = points
			self.points.changed = self.draw	#When points change, we redraw
			dim = len(shape(points.pos))

		if dim:
			if self.dim != dim:
				self.dim = dim
#				print "dim:", dim, "self.dim:", self.dim
				self.build_menu()
				self.ready = True
		else:
#			print "no points"
			return

#		if not (self.points and len(shape(self.points.pos)) == self.dim):	#I won't change dimension if they don't tell me!
#			return

#	def set_active(self, state):
#		self.active = state

	def build_menu(self):
		combo_model=ListStore(str)
		for item in drawers[self.dim]:
			combo_model.append([item])
		self.combo.set_model(combo_model)
		if self.dr_type in drawers[self.dim]:
			self.combo.set_active(drawers[self.dim].index(self.dr_type))
		else:
			self.combo.set_active(0)
#		print "built"

	def dump(self, dumpdir):
		print "dumping in", dumpdir
		self.dumpindex = 0

		if dumpdir:				#This is a "start" command
			self.dumpdir = dumpdir
			self.size = self.drawing.window.get_size()
			self.points.changed = self.dump_draw
		else:						#This is a "stop" command
			self.dumpdir = None
			self.size = None
			self.points.changed = self.draw
