#!/usr/bin/env python
import random
import gtk
from GTK import *
from GDK import *

def int2string(int, dims):
	str = ""
	for i in range(dims):
		if int & (1 << i): str = '1' + str
		else: str = '0' + str
	return str

def lg(power2):
	i = 0
	while power2 > 0:
		i = i + 1
		power2 = power2 >> 1
	return i - 1

def hamming_distance(int1, int2, dims):
	d = 0
	xor = int1 ^ int2
	for i in range(dims):
		if xor & (1 << i): d = d + 1
	return d



RecursiveArrays = {}

def build_recursive_array(ppd):
	RecursiveArrays[ppd] = []
	dims_on_2 = lg(ppd)
	dims = dims_on_2 * 2
	for int in xrange(1 << dims):
		new = 0
		for i in range(0, dims, 2):
			if int & (1 << i): new = new + (1 << (i/2))
			if int & (1 << (i+1)): new = new + (1 << (i/2+dims_on_2))
		x = new/ppd
		y = new%ppd
		RecursiveArrays[ppd].append((x,y))

def i2p_linear(int, ppd):
	return (int/ppd, int%ppd)

def p2i_linear(x, y, ppd):
	return x*ppd + y

def i2p_recursive(int, ppd):
	try:
		return RecursiveArrays[ppd][int]
	except KeyError:
		build_recursive_array(ppd)
		return i2p_recursive(int, ppd)

def p2i_recursive(x, y, ppd):
	try:
		return RecursiveArrays[ppd].index((x,y))
	except KeyError:
		build_recursive_array(ppd)



class FitnessFunction:

	def __init__(self, dims):
		self.dims = dims
		self.num_points = 1 << dims				# 2^dims
		self.sqrt_points = 1 << (dims / 2)		# sqrt(self.num_points)
		self.fitnesses = map(lambda i, fn=self.fitness: fn(i), 
							xrange(self.num_points))
		self.min_fitness = min(self.fitnesses)
		self.max_fitness = max(self.fitnesses)
		self.fitness = lambda i, f=self.fitnesses: f[i]

	#
	# Return fitness of solution corresponding to int.
	#
	def fitness(self, int):
		return -1

	#
	# Return list of points that are Hamming distance 1 away from int.
	#
	def neighbours(self, int):
		neighbours = []
		for i in range(self.dims):
			bit = 1 << i
			if bit & int: neighbours.append(int - bit)
			else: neighbours.append(int + bit)
		return neighbours

	# Return whether all points in neighbourhood have lesser fitness than n.
	def neighbours_are_worse(self, n, neighbourhood):
		return reduce(lambda x, y, ff=self.fitness, f=self.fitness(n):
					x and ff(y) <= f, neighbourhood, 1)

	#
	# Return list of points that are local optima.
	#
	def local_optima(self):
		optima = filter(lambda x, 
						nf=self.neighbours,
						ff=self.neighbours_are_worse: ff(x, nf(x)), 
						xrange(self.num_points))
		# Overwrite this function so that we don't have to call it again.
		self.local_optima = lambda opts=optima: opts
		return optima

	#
	# Return list of points that are global optima.
	#
	def optima(self):
		return filter(lambda x,fn=self.fitnesses,mf=self.max_fitness:
							fn[x] == mf, xrange(self.num_points))
		

	#
	# Recursive depth first search.
	#
	def dfs(self, int, seen, func):
		neighbours = filter(lambda n, s=seen, f=func, i=int:
								not n in s and f(i, n),
								self.neighbours(int))
		for n in neighbours:
			seen.append(n)
			self.dfs(n, seen, func)
		#seen = seen + neighbours
		#map(lambda n, s=seen, f=func, srch=self.dfs: srch(n, s, f), neighbours)
				#self.dfs(n, seen, func)

	def basin_func(self, current, next):
		return self.fitness(next) <= self.fitness(current)

	def optima_func(self, current, next):
		return self.fitness(next) >= self.fitness(current)

	#
	# Return list of points that are in the basin of attraction for int.
	# That is, all points which can be reached by successions of single-bit
	# changes that result in worse solutions; alternatively, the set of
	# points from which hillclimbing may find int.
	#
	def basin(self, int):
		basin = [ int ]
		self.dfs(int, basin, self.basin_func)
		return basin

	#
	# Return list of points that have higher fitness than int, and that may
	# be reached through hillclimbing single-bit mutations.  Opposite of
	# basin.
	#
	def reachable(self, int):
		reachable = [ int ]
		self.dfs(int, reachable, self.optima_func)
		return reachable

	def dfs2(self, int, seen):
		neighbours = self.neighbours(int)
		fitnesses = map(self.fitness, neighbours)
		m = max(fitnesses)
		if m <= self.fitness(int): return
		r = range(len(neighbours))
		best_neighbour_indices = filter(lambda i,
								f=fitnesses,
								max=m: f[i] >= max,
								r)
		n = neighbours[random.choice(best_neighbour_indices)]
		seen.append(n)
		self.dfs2(n, seen)

	def steepest_climb(self, int):
		path = [ int ]
		self.dfs2(int, path)
		return path

	def dfs3(self, int, seen):
		neighbours = self.neighbours(int)
		fitnesses = map(self.fitness, neighbours)
		fitness = self.fitness(int)
		if fitness >= max(fitnesses): return
		better_neighbour_indices = filter(lambda i,
								f=fitnesses,
								cutoff=fitness: f[i] > cutoff,
								range(len(neighbours)))
		n = neighbours[random.choice(better_neighbour_indices)]
		seen.append(n)
		self.dfs3(n, seen)

	def climb(self, int):
		path = [ int ]
		self.dfs3(int, path)
		return path

class JestersCap(FitnessFunction):
	eight = (
		32, 18, 18, 20, 18, 16, 16, 18, 18, 16, 16, 18, 20, 18, 18, 24,
		18, 12, 12, 14, 12, 10, 10, 12, 12, 10, 10, 12, 14, 12, 12, 18,
		18, 12, 12, 14, 12, 10, 10, 12, 12, 10, 10, 12, 14, 12, 12, 18,
		20, 14, 14, 16, 14, 12, 12, 14, 14, 12, 12, 14, 16, 14, 14, 20,
		18, 12, 12, 14, 12, 10, 10, 12, 12, 10, 10, 12, 14, 12, 12, 18,
		16, 10, 10, 12, 10,  8,  8, 10, 10,  8,  8, 10, 12, 10, 10, 16,
		16, 10, 10, 12, 10,  8,  8, 10, 10,  8,  8, 10, 12, 10, 10, 16,
		18, 12, 12, 14, 12, 10, 10, 12, 12, 10, 10, 12, 14, 12, 12, 18,
		18, 12, 12, 14, 12, 10, 10, 12, 12, 10, 10, 12, 14, 12, 12, 18,
		16, 10, 10, 12, 10,  8,  8, 10, 10,  8,  8, 10, 12, 10, 10, 16,
		16, 10, 10, 12, 10,  8,  8, 10, 10,  8,  8, 10, 12, 10, 10, 16,
		18, 12, 12, 14, 12, 10, 10, 12, 12, 10, 10, 12, 14, 12, 12, 18,
		20, 14, 14, 16, 14, 12, 12, 14, 14, 12, 12, 14, 16, 14, 14, 20,
		18, 12, 12, 14, 12, 10, 10, 12, 12, 10, 10, 12, 14, 12, 12, 18,
		18, 12, 12, 14, 12, 10, 10, 12, 12, 10, 10, 12, 14, 12, 12, 18,
		24, 18, 18, 20, 18, 16, 16, 18, 18, 16, 16, 18, 20, 18, 18, 32,
	)

	#
	# Return fitness of solution corresponding to int.
	#
	def fitness(self, int):
		if self.dims == 8:
			return JestersCap.eight[int]
		elif self.dims == 16:
			hi = int >> 8
			lo = int & 0xFF
			score = JestersCap.eight[hi] + JestersCap.eight[lo]
			if hi == 0xFF or hi == 0: score = score + 8
			if lo == 0xFF or lo == 0: score = score + 8
			return score
		else: return 0


class Random(FitnessFunction):

	def fitness(self, i):
		return int(random.random() * 100)

NN_1 = Random



class N0(FitnessFunction):

	def __init__(self, dims):
		val0 = map(lambda i: int(random.random() * 1000), range(dims))
		val1 = map(lambda i: int(random.random() * 1000), range(dims))
		self.val0 = map(lambda i, l0=val0, l1=val1: min( (l0[i],l1[i]) ), 
			range(dims))
		self.val1 = map(lambda i, l0=val0, l1=val1: max( (l0[i],l1[i]) ), 
			range(dims))
		FitnessFunction.__init__(self, dims)

	def fitness(self, int):
		f = 0
		for b in range(self.dims):
			bit = 1 << b
			if int & bit: f = f + self.val1[b]
			else: f = f + self.val0[b]
		return f


class N2(FitnessFunction):

	def __init__(self, dims):
		self.vals = {}
		for center in range(dims):
			left = (center + 1) % dims
			right = (center + dims - 1) % dims
			for i in range(8):
				b3, b2, b1 = i & 1, (i & 2) >> 1, (i & 4) >> 2
				v = (b1 << left) + (b2 << center) + (b3 << right)
				self.vals[(v, center+left+right)] = int(random.random() * 150)
		FitnessFunction.__init__(self, dims)

	def fitness(self, int):
		f = 0
		for center in range(self.dims):
			left = (center + 1) % self.dims
			right = (center + self.dims - 1) % self.dims
			if int & (1 << right): b3 = 1
			else: b3 = 0
			if int & (1 << center): b2 = 1
			else: b2 = 0
			if int & (1 << left): b1 = 1
			else: b1 = 0
			v = (b1 << left) + (b2 << center) + (b3 << right)
			f = f + self.vals[ (v, center+left+right) ]
		return f



class ColourHandler:

	def __init__(self, display):
		self.display = display
		self.cid = self.display.canvas.connect('configure_event', self.create_colours)

	def create_colours(self, *args):
		self.fg = self.display.canvas.get_style().black_gc.foreground
		self.greys = {}
		self.greens = {}
		self.reds = {}
		self.oranges = {}
		cmap = self.display.canvas.get_colormap()
		for i in xrange(self.display.model.num_points):
			fitness = self.display.model.fitness(i)
			try:
				self.greys[fitness]
				self.greens[fitness]
				self.reds[fitness]
				self.oranges[fitness]
			except KeyError:
				lvl = (fitness - self.display.model.min_fitness) * 65535 \
					/ (self.display.model.max_fitness - self.display.model.min_fitness)
				self.greys[fitness] = cmap.alloc(lvl, lvl, lvl)
				self.reds[fitness] = cmap.alloc(0xFFFF,lvl, lvl)
				self.greens[fitness] = cmap.alloc(0xFFFF-lvl, 0xFFFF, 0xFFFF-lvl)
				self.oranges[fitness] = cmap.alloc(0xFFFF-lvl, lvl, 0)
				self.blue = cmap.alloc(0, 0, 0xFFFF)
		try:
			n = self.cid
			del(self.cid)
			self.display.canvas.disconnect(n) # Created colours, ignore configures.
		except AttributeError: pass

	def grey(self, int):
		return self.greys[self.display.model.fitness(int)]

	def orange(self, int):
		return self.oranges[self.display.model.fitness(int)]

	def red(self, int):
		return self.reds[self.display.model.fitness(int)]

	def green(self, int):
		return self.greens[self.display.model.fitness(int)]
	

class Display(gtk.GtkWindow):

	def __init__(self, model):
		self.model = model
		gtk.GtkWindow.__init__(self, type=WINDOW_TOPLEVEL, title='Hyperspace Graph Paper')
		self.topLevelLayout = gtk.GtkVBox()
		self.topLevelLayout.show()
		self.add(self.topLevelLayout)
		self.canvas = gtk.GtkDrawingArea()
		self.canvas.set_usize(600, 600)
		self.canvas.connect('expose_event', self.refresh)
		self.canvas.connect_after('configure_event', self.configure_canvas)
		self.colours = ColourHandler(self)
		self.canvas.add_events(BUTTON_PRESS_MASK)
		self.canvas.connect('button_press_event', self.button_press)
		self.canvas.show()
		self.topLevelLayout.pack_start(self.canvas, expand=TRUE, fill=TRUE)
		sep = gtk.GtkHSeparator()
		sep.show()
		self.topLevelLayout.pack_start(sep, expand=FALSE, fill=FALSE)
		# What to highlight
		toolbar = gtk.GtkHBox()
		toolbar.show()
		self.topLevelLayout.pack_start(toolbar, expand=FALSE, fill=TRUE)
		tool_group = None
		for tool in [ 'None', 'Neighbours', 'Basin',
				'Peaks', 'Optima', 'Steepest', 'Any', 'ROptima' ]:
			r = gtk.GtkRadioButton(tool_group, tool)
			r.set_data('highlight', tool)
			r.connect('toggled', self.highlight_selected)
			tool_group = r
			toolbar.pack_start(r, expand=FALSE, fill=FALSE)
			r.show()
		# Bits to highlight
		"""
		subtoolbar = gtk.GtkHBox()
		subtoolbar.show()
		self.topLevelLayout.pack_start(subtoolbar, expand=FALSE, fill=TRUE)
		for bit in range(8):
			r = gtk.GtkRadioButton(tool_group, `bit`)
			r.set_data('bit', bit)
			r.connect('toggled', self.highlight_bits_selected)
			tool_group = r
			subtoolbar.pack_start(r, expand=FALSE, fill=FALSE)
			r.show()
		sep = gtk.GtkHSeparator()
		sep.show()
		self.topLevelLayout.pack_start(sep, expand=FALSE, fill=FALSE)
		"""
		# What function to show
		functionbar = gtk.GtkHBox()
		functionbar.show()
		self.topLevelLayout.pack_start(functionbar, expand=FALSE, fill=TRUE)
		f_group = None
		for function in [ 'JC8', 'JC16', 'N8K7', 'N8K0', 'N8K2']:
			r = gtk.GtkRadioButton(f_group, function)
			r.set_data('function', function)
			r.connect('toggled', self.function_selected)
			f_group = r
			functionbar.pack_start(r, expand=FALSE, fill=FALSE)
			r.show()
		sep = gtk.GtkHSeparator()
		sep.show()
		self.topLevelLayout.pack_start(sep, expand=FALSE, fill=FALSE)
		# How to map points to ints and vice-versa
		self.i2p = i2p_recursive
		self.p2i = p2i_recursive
		transformbar = gtk.GtkHBox()
		transformbar.show()
		self.topLevelLayout.pack_start(transformbar, expand=FALSE, fill=TRUE)
		tgroup = None
		for transform in [ 'Recursive', 'Linear' ]:
			r = gtk.GtkRadioButton(tgroup, transform)
			r.set_data('transform', transform)
			r.connect('toggled', self.transform_selected)
			tgroup = r
			transformbar.pack_start(r, expand=FALSE, fill=FALSE)
			r.show()
		"""
		self.printButton = gtk.GtkButton("Print")
		self.printButton.connect('clicked', self.printmodel)
		self.printButton.show()
		self.topLevelLayout.pack_start(self.printButton, expand=FALSE, fill=TRUE)
		"""
		# Statusbar
		self.status = gtk.GtkStatusbar()
		self.status.show()
		self.topLevelLayout.pack_start(self.status, expand=FALSE, fill=TRUE)
		# blah
		self.connect('delete_event', gtk.mainquit)
		self.connect('destroy_event', gtk.mainquit)
		self.highlight = []
		self.selected = None
		self.highlight_style = 'None'

	def printmodel(self, *args):
		print 'optima: %s' % `self.model.optima()`
		optimum = self.model.optima()[0]
		print 'distance from optimum(%d), fitness' % optimum
		for i in range(self.model.num_points):
			d = hamming_distance(optimum, i, self.model.dims)
			f = self.model.fitnesses[i]
			print "%d, %d" % (d, f)
		#print self.model.fitnesses

	def set_status(self, msg):
		try: id = self.statusid
		except: id = self.statusid = self.status.get_context_id('fubar')
		self.status.push(id, msg)

	def configure_canvas(self, *args):
		self.width = self.canvas.get_allocation()[2]
		self.height = self.canvas.get_allocation()[3]
		self.point_width = int(self.width / self.model.sqrt_points)
		self.point_height = int(self.height / self.model.sqrt_points)
		self.base_pixmap = gtk.create_pixmap(self.canvas, self.width, self.height)
		self.pixmap = gtk.create_pixmap(self.canvas, self.width, self.height)
		self.canvas_gc = self.canvas.get_parent_window().new_gc()
		print `self.canvas_gc`
		self.draw_basic()
		self.draw_highlights()
		print "points are: %d, %d" % (self.point_width, self.point_height)

	def transform_selected(self, widget):
		if widget.active:
			t = widget.get_data('transform')
			if t == 'Recursive':
				self.i2p = i2p_recursive
				self.p2i = p2i_recursive
			elif t == 'Linear':
				self.i2p = i2p_linear
				self.p2i = p2i_linear
			self.draw_basic()
			self.update_highlight_group()

	def function_selected(self, widget):
		if widget.active:
			old_dims = self.model.dims
			fn = widget.get_data('function')
			if fn == 'JC8':
				self.model = JestersCap(8)
			if fn == 'JC16':
				self.model = JestersCap(16)
			elif fn == 'N8K7':
				self.model = NN_1(8)
			elif fn == 'N8K0':
				self.model = N0(8)
			elif fn == 'N8K2':
				self.model = N2(8)
			self.colours.create_colours()
			if self.model.dims != old_dims: self.configure_canvas()
			else:
				self.draw_basic()
				self.update_highlight_group()

	def update_highlight_group(self):
		if self.highlight_style == 'Optima':
			self.highlight = self.model.local_optima()
			self.selected = None
		elif self.highlight_style == 'Bits':
			self.selected = None
			self.highlight = filter(lambda i, b=self.hbit: i & (1 << b), 
										range(self.model.num_points))
		elif self.selected == None or self.highlight_style == 'None':
			self.selected = None
			self.highlight = []
		elif self.highlight_style == 'Basin':
			self.highlight = self.model.basin(self.selected)
		elif self.highlight_style == 'Neighbours':
			self.highlight = self.model.neighbours(self.selected)
		elif self.highlight_style == 'Peaks':
			self.highlight = self.model.reachable(self.selected)
		elif self.highlight_style == 'Steepest':
			self.highlight = self.model.steepest_climb(self.selected)
		elif self.highlight_style == 'Any':
			self.highlight = self.model.climb(self.selected)
		elif self.highlight_style == 'ROptima':
			optima = self.model.local_optima()
			reachable = self.model.reachable(self.selected)
			self.highlight = filter(lambda i, r=reachable: i in r, optima)
		self.draw_highlights()

	def highlight_selected(self, widget):
		if widget.active:
			self.highlight_style = widget.get_data('highlight')
			self.update_highlight_group()

	def highlight_bits_selected(self, widget):
		if widget.active:
			self.hbit = widget.get_data('bit')
			self.highlight_style = 'Bits'
			self.update_highlight_group()

	def button_press(self, widget, event):
		x = int(event.x) / self.point_width
		y = int(event.y) / self.point_height
		n = self.p2i(x, y, self.model.sqrt_points)
		self.selected = n
		self.set_status('%s: fitness = %d.  %d highlighted' %
			(int2string(n, self.model.dims), self.model.fitness(n), len(self.highlight)))
		self.update_highlight_group()

	def refresh(self, *args):
		self.canvas.draw_pixmap(self.canvas.get_style().fg_gc[1],
			self.pixmap,
			0, 0, 0, 0,
			self.width, self.height)

	def draw_basic(self):
		style = self.canvas.get_style()
		map(self.draw_point, xrange(self.model.num_points))
		self.canvas_gc.foreground = self.colours.fg
	
	def draw_point(self, i):
		x, y = self.i2p(i, self.model.sqrt_points)
		x = x * self.point_width
		y = y * self.point_height
		self.canvas_gc.foreground = self.colours.grey(i)
		gtk.draw_rectangle(self.base_pixmap, self.canvas_gc, TRUE, x, y, 
				self.point_width, self.point_height)

	def draw_hpoint(self, i, target_fitness):
		x, y = self.i2p(i, self.model.sqrt_points)
		x = x * self.point_width
		y = y * self.point_height
		fitness = self.model.fitness(i)
		if fitness > target_fitness:
			col = self.colours.green(i)
		elif fitness < target_fitness:
			col = self.colours.red(i)
		else:
			col = self.colours.orange(i)
		self.canvas_gc.foreground = col
		gtk.draw_rectangle(self.pixmap, self.canvas_gc, TRUE, x, y, 
				self.point_width, self.point_height)
		self.canvas_gc.foreground = self.colours.blue
		gtk.draw_line(self.pixmap, self.canvas_gc, x, y, 
									x+self.point_width, y+self.point_height)
		gtk.draw_line(self.pixmap, self.canvas_gc, x+self.point_width, y, 
									x, y+self.point_height)


	def draw_highlights(self, *args):
		gtk.draw_pixmap(self.pixmap, 
			self.canvas_gc,
			self.base_pixmap,
			0, 0, 0, 0,
			self.width, self.height)
		style = self.canvas.get_style()
		if self.selected != None:
			target_fitness = self.model.fitness(self.selected)
		else: target_fitness = 0
		map(lambda i,
				t=target_fitness,
				fn=self.draw_hpoint,
				s=style: fn(i, t),
				self.highlight)
		gc = style.black_gc
		pd = self.model.sqrt_points
		for i in range(self.model.sqrt_points+1):
			if i & 1:
				gtk.draw_line(self.pixmap, gc, 0, i*self.point_height, 
							self.width, i*self.point_height)
				gtk.draw_line(self.pixmap, gc, i*self.point_width, 0,
							i*self.point_width, self.height)
			elif i & 2:
				gtk.draw_rectangle(self.pixmap, gc, TRUE, 0, i*self.point_height-1, 
							self.width, 3)
				gtk.draw_rectangle(self.pixmap, gc, TRUE, i*self.point_width-1, 0,
							3, self.height)
			elif i & 4:
				gtk.draw_rectangle(self.pixmap, gc, TRUE, 0, i*self.point_height-2, 
							self.width, 5)
				gtk.draw_rectangle(self.pixmap, gc, TRUE, i*self.point_width-2, 0,
							5, self.height)
			elif i & 8:
				gtk.draw_rectangle(self.pixmap, gc, TRUE, 0, i*self.point_height-3, 
							self.width, 7)
				gtk.draw_rectangle(self.pixmap, gc, TRUE, i*self.point_width-3, 0,
							7, self.height)
			elif i & 16:
				gtk.draw_rectangle(self.pixmap, gc, TRUE, 0, i*self.point_height-4, 
							self.width, 9)
				gtk.draw_rectangle(self.pixmap, gc, TRUE, i*self.point_width-4, 0,
							9, self.height)
		self.canvas_gc.foreground = self.colours.fg
		self.refresh()


j = JestersCap(8)
d = Display(j)
d.show()
gtk.mainloop()