""" Quicksort visualization applet.

My first attempt to write an applet and my second real python program. :-)
"""

from Tkinter import *
import string

#stack= list, push = insert(0,_), pop = remove(0)

class SortState:
    """Internal state of quick sort parition."""

    def __init__(self, left_bound=-1, right_bound=-1,
		 pivot=-1, left=-1, right=-1):
	self.left_bound = left_bound
	self.right_bound = right_bound
	self.pivot = pivot
	self.left = left
	self.right = right
	
class Quicksort:
    """Stack machine that implementes quicksort.

    The method 'next' performs a machine operation, which modifies
    the current state. The ops 'start' and 'done' modify the stack too.
    Each operation also sets the instruction pointer as appropriate.

    See Knuth or Cormen, Leiserson, & Rivest for description of
    quicksort algorithm.
    """
	
    def __init__(self, parent, array):
	self.ui = parent
	self.array = array
	self.stack = [SortState(0, len(array)-1)]
	self.state = SortState()
	self.ip = self.start
	self.op_name = { self.start : 'Start',
			 self.slide_left : 'Slide left',
			 self.slide_right : 'Slide right',
			 self.exchange_or_done : 'Exchange or Done',
			 self.done : 'Done',
			 self.noop: '' }
	
    def next(self):
	next_op = self.ip()
	self.ui.state_label.config(text=self.op_name[next_op])
	self.ip = next_op

   #From here down are machine operations

    def noop(self):
	return self.noop

    def start(self):
	self.ui.reset_bar_color('active')
	if len(self.stack) == 0:
	    return self.noop

	# pop the next partition off the stack
	# and intialize current state with it
	cur = self.stack[0]
	del self.stack[0]
	cur.pivot = self.array[cur.left_bound]
	cur.left = cur.left_bound - 1
	cur.right = cur.right_bound + 1
	self.state = cur
	self.ui.incr_comparisons()
	if cur.left_bound < cur.right_bound:
	    for bar in range(cur.left_bound+1, cur.right_bound+1):
		self.ui.set_bar_color(bar, 'active')
	    self.ui.set_bar_color(cur.left_bound, 'pivot')
	    return self.slide_left
	else:
	    return self.start

    def slide_left(self):
	state = self.state
	state.right = state.right - 1
	while self.array[state.right] > state.pivot:
	    self.ui.incr_comparisons()
	    state.right = state.right - 1
	self.ui.incr_comparisons()
	self.ui.set_bar_color(state.right, 'left')
	return self.slide_right

    def slide_right(self):
	state = self.state
	state.left = state.left + 1
	while self.array[state.left] < state.pivot:
	    self.ui.incr_comparisons()
	    state.left = state.left + 1
	self.ui.incr_comparisons()
	self.ui.set_bar_color(state.left, 'right')
	return self.exchange_or_done

    def exchange_or_done(self):
	state = self.state
	self.ui.reset_bar_color('left')
	self.ui.reset_bar_color('right')
	if state.left < state.right:
	    self.ui.incr_comparisons()
	    self.ui.swap(state.left,state.right)
	    return self.slide_left
	else:
	    return self.done

    def done(self):
	state = self.state
	self.ui.reset_bar_color('pivot')
	if state.left_bound < state.right:
	    self.stack.insert(0, SortState(state.left_bound, state.right,
					   -1, -1, -1))
	if state.right + 1 < state.right_bound:
	    self.stack.insert(0, SortState(state.right + 1, state.right_bound,
					   -1, -1, -1))
	return self.start

class QuicksortUI:
    """Tk interface for sort.
    """

    def __init__(self, parent, bar_height=80, bar_width=10, bar_gap=2):
	self.parent = parent
		
	self.sort = Quicksort(self,[])
	self.bars = []
	self.exchanges = 0
	self.comparisons = 0
	self.interrupt = 0
	self.ready = 0
	
	self.prev_bar_color = {}
	self.colormap = {'default' : 'red',
			 'pivot' : 'green',
			 'left' : 'yellow',
			 'right' : 'yellow',
			 'active' : 'white'}
	
	self.colored_bars = {'pivot' : [],
			     'bound' : [],
			     'left' : [],
			     'right' : [],
			     'active' : []}
	
	self.bar_height = bar_height
	self.bar_width = bar_width
	self.bar_gap = bar_gap
	
	# setup the buttons and things
	
	f = Frame(parent)
	t = Label(f, text="Enter list of numbers")
	v = StringVar()
	v.set("1 10 9 2 3 8 7 4 6 5")
	self.array_entry = Entry(f, width=40, textvariable=v)
	t.pack(side=LEFT)
	self.array_entry.pack(side=LEFT)
	
	reset = Button(parent,
		       text = 'Reset sort display',
		       command = self.read_array)
	f.pack()
	reset.pack(pady=10)
	
	self.canvas = Canvas(parent,
				  height = bar_height + 2 * bar_gap,
				  width = (bar_width + bar_gap) * 20,
				  relief = SUNKEN,
				  borderwidth = 2)
	self.canvas.pack()
	
	f = Frame(parent)
	self.l_exch = Label(f,
			    text = "Exchanges: %d" % self.exchanges,
			    width = 15)
	self.l_comp = Label(f,
			    text="Comparisons: %d" % self.comparisons,
			    width = 15)
	self.l_exch.pack(side=LEFT,anchor="w")
	self.l_comp.pack(side=RIGHT,anchor="e")
	f.pack(fill="x")
	
	f = Frame(parent)
	t = Label(f, text="Operation:")
	self.state_label = Label(f,
				 text = "")
	t.pack(side=LEFT)
	self.state_label.pack(side=LEFT)
	f.pack()
	
	self.sort_next = Button(parent,
				text='Single step',
				command=self.next_step)
	self.sort_next.pack(side=LEFT)
	
	self.sort_go = Button(parent,
			      text='Continuous',
			      command=self.go)
	self.sort_go.pack(side=RIGHT)
	
	self.sort_stop = Button(parent,
				text='Stop',
				command=self.stop)
	self.sort_stop.pack(side=RIGHT)
	self.read_array()


    def read_array(self):
	"""Read numbers from the entry item.

	No error checking performed currently."""

	self.sort = Quicksort(self, string.split(self.array_entry.get()))
	for i in range(len(self.sort.array)):
	    self.sort.array[i] = string.atoi (self.sort.array[i])
	self.initialize_canvas()
	self.set_exchanges(0)
	self.set_comparisons(0)
	self.ready = 1

    def initialize_canvas(self):
	"""Create rectangles to represent items being sorted.
	
	There is a bar for each number being sorted. The bar's
	height is proportional to the number and the color changes
	depending on what Quicksort is doing with the number.
	
	Because we can't use just number to identify a bar, use
	an array that is kept in the same order as the numbers 
	being sorted.
	"""
	self.canvas.delete('all')
	self.bars = []
	self.setup_translate(self.sort.array)
	self.canvas.config(width=len(self.sort.array)*(self.bar_width+self.bar_gap))
	for i in range(len(self.sort.array)):
	    x = 2 + i * (self.bar_width + self.bar_gap) \
		+ self.bar_gap
	    y = self.translate_height(self.sort.array[i])
	    bar = self.canvas.create_rectangle(x, y, x + self.bar_width,
					       self.translate_height(0),
					       fill=self.colormap['default'])
	    self.bars.insert(i, bar)
	    self.prev_bar_color[bar] = []
	    
	    # translate from user-enter numbers to actual canvas height
    def setup_translate(self,array):
	self.max_bar = max(array)
	min_bar = min(array)
	if min_bar > 0:
	    self.min_bar = 0
	else:
	    self.min_bar = min_bar
	self.offset = 1 + string.atoi(self.canvas['borderwidth']) * 2
	self.range = self.max_bar - self.min_bar
	self.top = self.bar_height + self.offset
	
    def translate_height(self,h):
	return self.top - self.bar_height * (h - self.min_bar) / self.range
		
    # three routines to control the exchanges label
    def set_exchanges(self, val):
	self.exchanges = val
	self.show_exchanges()

    def incr_exchanges(self):
	self.exchanges = self.exchanges + 1
	self.show_exchanges()

    def show_exchanges(self):
	self.l_exch.config(text="Exchanges: %d" % self.exchanges)

    # three routines to control the comparisons label
    def set_comparisons(self, val):
	self.comparisons = val
	self.show_comparisons()

    def incr_comparisons(self):
	self.comparisons = self.comparisons + 1
	self.show_comparisons()

    def show_comparisons(self):
	self.l_comp.config(text="Comparisons: %d" % self.comparisons)

    def swap(self,i,j):
	# a little complicated becauser we're swapping bars and
	# array slots
	num = self.sort.array[i]
	self.sort.array[i] = self.sort.array[j]
	self.sort.array[j] = num

	bar = self.bars[i]
	self.bars[i] = self.bars[j]
	self.bars[j] = bar

	self.move_bar(self.bars[i], j, i)
	self.move_bar(self.bars[j], i, j)

	self.incr_exchanges()
	
    def move_bar(self, bar, old_index, new_index):
	self.canvas.move (bar, self.bar_index_to_offset(new_index)
			  - self.bar_index_to_offset(old_index), 0)

	# routines for changing around the bar colors based on the
	# role it is playing in the current sort state.
	# when a bar's color is changed, the previous value is pushed
	# onto a stack in the prev_bar_color dict.
	#
	# ugly & messy.

    def get_bar_color(self,bar):
	color = self.canvas.itemconfig(bar, 'fill')[-1:]
	return color

    def set_bar_color(self,index,type):
	self.colored_bars[type].append(self.bars[index])
	self.prev_bar_color[self.bars[index]
			    ].insert(0,self.get_bar_color(self.bars[index]))
	self.canvas.itemconfig(self.bars[index],
			       fill=(self.colormap[type]))

    def reset_bar_color(self,type):
	for bar in self.colored_bars[type]:
	    oldcolor = self.prev_bar_color[bar][0]
	    del self.prev_bar_color[bar][0]
	    self.canvas.itemconfig(bar, fill=(oldcolor))
	self.colored_bars[type] = []

    def bar_index_to_offset(self,index):
	x = index * (self.bar_width+self.bar_gap) + self.bar_gap
	return x

    # methods bound to the buttons
    def next_step(self):
	if self.ready:
	    self.sort.next()

    def stop(self):
	if self.sort.ip != self.sort.noop:
	    self.interrupt = 1

    def go(self):
	if self.ready:
	    if self.interrupt == 1:
		self.interrupt = 0
		return
	    self.sort.next()
	    if self.sort.ip != self.sort.noop:
		self.parent.update_idletasks()
		self.parent.after(200, self.go)
		# have a control for the delay?

if __name__ == '__main__': # Are we run as a main program?
    root = Tk()
    app = QuicksortUI(root)
    root.mainloop()