# A simple GUI example using Tk frames
import matplotlib
matplotlib.use('TkAgg')
from matplotlib.numerix import arange, sin, pi
from matplotlib.axes import Subplot
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2TkAgg
from matplotlib.figure import Figure
import sys
from Numeric import *
from Agent import *
from visual import vector
from Tkinter import *
	
def start():
	#Global flag to start plotting on the first press of the button
	global firstRun
	global GO
	GO = TRUE
	if firstRun:
		firstRun = FALSE
		runProgram()
		
#Stop the simulation	
def stop():
	global GO
	global xValsX,xValsY,yValsX,yValsY
	#Clear the variable arrays so when Start is pressed again, it doesn't plot the old values twice
	if size(xValsX) != 0:
		n = size(xValsX)
		xValsX = [xValsX[n-1]]
		xValsY = [xValsY[n-1]]
		yValsX = [yValsX[n-1]]
		yValsY = [yValsY[n-1]]	
	GO = FALSE

#Project data onto plane from 3 space
def TwoDProjection(x):
	projectionMatrix = array([[ -1./2.,1./2.,0.],
				  [ 0,0,sqrt(3.)/2.]])
	return dot(projectionMatrix,x)

#Set the slider to match the value typed into the Entry box
def __alphaXupdateSlider(event):
	alphaXScale.set(parameterVal.get())
	
def __alphaYupdateSlider(event):
	alphaYScale.set(parameterVal.get())

def __betaXupdateSlider(event):
	betaXScale.set(parameterVal.get())

def __betaYupdateSlider(event):
	betaYScale.set(parameterVal.get())

#Set the value in the Entry. One for each parameter.
def alphaXUpdate(val):
	alphaX.set(val)
	
def alphaYUpdate(val):
	alphaY.set(val)
	
def betaXUpdate(val):
	betaX.set(val)

def betaYUpdate(val):
	betaY.set(val)

def runProgram():
	# Set initial conditions.
	u = -log(vector(xIC1.get(),yIC1.get(),zIC1.get()))
	v = -log(vector(xIC2.get(),yIC2.get(),zIC2.get()))
	
	x = vector(0,0,0)
	y = vector(0,0,0)
	
	while TRUE:
		master.update()
		count = 0				
		alpha = [alphaX.get(),alphaY.get()]
		beta = [betaX.get(),betaY.get()]
				
		while GO:
			master.update()
			if count == 100:
				global xValsX,xValsY,yValsX,yValsY
				n = size(xValsX)
				xValsX = [xValsX[n-1]]
				xValsY = [xValsY[n-1]]
				yValsX = [yValsX[n-1]]
				yValsY = [yValsY[n-1]]
				count = 0
			count += 1
			canvas.show()
			X = Agent(alpha[0],beta[0],rewardX)
			Y = Agent(alpha[1],beta[1],rewardY)
			u = X.RK4(y,u,dt)
			v = Y.RK4(x,v,dt)
			x = exp(-u)/sum(exp(-u))
			y = exp(-v)/sum(exp(-v))
			XTrans = array(TwoDProjection(x))
			YTrans = array(TwoDProjection(y))
			xValsX.append(XTrans[0])
			xValsY.append(XTrans[1])
			yValsX.append(YTrans[0]+1)
			yValsY.append(YTrans[1])
			graph.plot(xValsX,xValsY,'b')
			graph.plot(yValsX,yValsY,'g')
			alpha = [alphaX.get(),alphaY.get()]
			beta = [betaX.get(),betaY.get()]

#END FUNCTION DEFINITIONS

firstRun = TRUE
GO = FALSE
#Build GUI environment
master = Tk()
master.title("Two Agent MADS Controls")
#Figure widget for holding the graph
figure = Figure(figsize=(5,4), dpi=100)
graph = figure.add_subplot(111)
canvas = FigureCanvasTkAgg(figure,master)
#canvas.show()
canvas._tkcanvas.pack(side=TOP, fill=BOTH, expand=1)

#Frame to hold Start,Stop, and Clear buttons
bFrame = Frame(master)
bFrame.pack(side=LEFT)
#Create start button
bStart = Button(bFrame,text='Start', command=start)
bStart.pack(side=LEFT)
#Create stop button
bStop = Button(bFrame,text='Stop', command=stop)
bStop.pack(side=LEFT)

#Create sliders to control values of alpha and beta, for both agents.
#Link the values of the sliders to a text box, so the user can manually enter values to use, for better precision than the slider alone.

#Slider for alphaX
#Create frame to hold slider and text box with appropriate label
alphaXFrame = Frame(master)
alphaXFrame.pack(side=TOP)
alphaXLabel = Label(alphaXFrame,text='alphaX')
alphaXLabel.pack(side=LEFT)
alphaXScale = Scale(alphaXFrame,from_=0.0,to=.9999,orient=HORIZONTAL,resolution=.0001,showvalue=FALSE,command=alphaXUpdate)
alphaXScale.pack(side=LEFT)
alphaXScale.set(0.0198)
alphaX = DoubleVar()
alphaXentry_field = Entry(alphaXFrame,width=6,textvariable=alphaX)
alphaXentry_field.pack(side=LEFT)
alphaXentry_field.event_add("<<ValueEntered>>", "<KeyPress-0>", "<KeyPress-1>", "<KeyPress-2>", "<KeyPress-3>", "<KeyPress-4>", "<KeyPress-5>", "<KeyPress-6>", "<KeyPress-7>", "<KeyPress-8>", "<KeyPress-9>" )
alphaXentry_field.bind("<<ValueEntered>>",__alphaXupdateSlider)

#Create slider for alphaY
alphaYFrame = Frame(master)
alphaYFrame.pack(side=TOP)
alphaYLabel = Label(alphaYFrame,text='alphaY')
alphaYLabel.pack(side=LEFT)
alphaYScale = Scale(alphaYFrame,from_=0.0,to=.9999,orient=HORIZONTAL,resolution=.0001,showvalue=FALSE,command=alphaYUpdate)
alphaYScale.pack(side=LEFT)
alphaYScale.set(0.01)
alphaY = DoubleVar()
alphaYentry_field = Entry(alphaYFrame,width=6,textvariable=alphaY)
alphaYentry_field.pack(side=LEFT)
alphaYentry_field.event_add("<<ValueEntered>>", "<KeyPress-0>", "<KeyPress-1>", "<KeyPress-2>", "<KeyPress-3>", "<KeyPress-4>", "<KeyPress-5>", "<KeyPress-6>", "<KeyPress-7>", "<KeyPress-8>", "<KeyPress-9>" )
alphaYentry_field.bind("<<ValueEntered>>",__alphaYupdateSlider)

#Create slider for betaX
betaXFrame = Frame(master)
betaXFrame.pack(side=TOP)
betaXLabel = Label(betaXFrame,text='betaX')
betaXLabel.pack(side=LEFT)
betaXScale = Scale(betaXFrame,from_=0.0,to=20.0,orient=HORIZONTAL,resolution=.1,showvalue=FALSE,command=betaXUpdate)
betaXScale.pack(side=LEFT)
betaXScale.set(1.0)
betaX = DoubleVar()
betaXentry_field = Entry(betaXFrame,width=6,textvariable=betaX)
betaXentry_field.pack(side=LEFT)
betaXentry_field.event_add("<<ValueEntered>>", "<KeyPress-0>", "<KeyPress-1>", "<KeyPress-2>", "<KeyPress-3>", "<KeyPress-4>", "<KeyPress-5>", "<KeyPress-6>", "<KeyPress-7>", "<KeyPress-8>", "<KeyPress-9>" )
betaXentry_field.bind("<<ValueEntered>>",__betaXupdateSlider)

#Create slider for betaY
betaYFrame = Frame(master)
betaYFrame.pack(side=TOP)
betaYLabel = Label(betaYFrame,text='betaY')
betaYLabel.pack(side=LEFT)
betaYScale = Scale(betaYFrame,from_=0.0,to=20.0,orient=HORIZONTAL,resolution=.1,showvalue=FALSE,command=betaYUpdate)
betaYScale.pack(side=LEFT)
betaYScale.set(2.0)
betaY = DoubleVar()
betaYentry_field = Entry(betaYFrame,width=6,textvariable=betaY)
betaYentry_field.pack(side=LEFT)
betaYentry_field.event_add("<<ValueEntered>>", "<KeyPress-0>", "<KeyPress-1>", "<KeyPress-2>", "<KeyPress-3>", "<KeyPress-4>", "<KeyPress-5>", "<KeyPress-6>", "<KeyPress-7>", "<KeyPress-8>", "<KeyPress-9>" )
betaYentry_field.bind("<<ValueEntered>>",__betaYupdateSlider)

#GO = FALSE
icFrame = Frame(master)
icFrame.pack(side=LEFT)
icLabel = Label(icFrame,text='Initial Conditions')
icLabel.pack(side=TOP)


xIC1 = DoubleVar()
yIC1 = DoubleVar()
zIC1 = DoubleVar()
xIC1.set(.2)
yIC1.set(.5)
zIC1.set(.3)

xIC2 = DoubleVar()
yIC2 = DoubleVar()
zIC2 = DoubleVar()
xIC2.set(.1)
yIC2.set(.5)
zIC2.set(.4)


agent1Frame = Frame(icFrame)
agent1Frame.pack(side=LEFT)
agent2Frame = Frame(icFrame)
agent2Frame.pack(side=RIGHT)

agentLabel1 = Label(agent1Frame,text='Agent 1')
agentLabel1.pack(side=TOP)
agentLabel2 = Label(agent2Frame,text='Agent 2')
agentLabel2.pack(side=TOP)

icXFrame1 = Frame(agent1Frame)
icXFrame1.pack(side=TOP)
icYFrame1 = Frame(agent1Frame)
icYFrame1.pack(side=TOP)
icZFrame1 = Frame(agent1Frame)
icZFrame1.pack(side=TOP)

icXLabel1 = Label(icXFrame1,text='x=')
icXLabel1.pack(side=LEFT)
icXField1 = Entry(icXFrame1,width=5,textvariable=xIC1)
icXField1.pack(side=RIGHT)

icYLabel1 = Label(icYFrame1,text='y=')
icYLabel1.pack(side=LEFT)
icYField1 = Entry(icYFrame1,width=5,textvariable=yIC1)
icYField1.pack(side=RIGHT)

icZLabel1 = Label(icZFrame1,text='z=')
icZLabel1.pack(side=LEFT)
icZField1 = Entry(icZFrame1,width=5,textvariable=zIC1)
icZField1.pack(side=RIGHT)

icXFrame2 = Frame(agent2Frame)
icXFrame2.pack(side=TOP)
icYFrame2 = Frame(agent2Frame)
icYFrame2.pack(side=TOP)
icZFrame2 = Frame(agent2Frame)
icZFrame2.pack(side=TOP)

icXLabel2 = Label(icXFrame2,text='x=')
icXLabel2.pack(side=LEFT)
icXField2 = Entry(icXFrame2,width=5,textvariable=xIC2)
icXField2.pack(side=RIGHT)

icYLabel2 = Label(icYFrame2,text='y=')
icYLabel2.pack(side=LEFT)
icYField2 = Entry(icYFrame2,width=5,textvariable=yIC2)
icYField2.pack(side=RIGHT)

icZLabel2 = Label(icZFrame2,text='z=')
icZLabel2.pack(side=LEFT)
icZField2 = Entry(icZFrame2,width=5,textvariable=zIC2)
icZField2.pack(side=RIGHT)

#END GUI CONFIGURATION

#Step size
dt = .1

# Tiebreak parameter for reward matrix (between -1 and 1)
e = [0.5,-0.3]

#reward matrix for agent 1
rewardX = array([[2./3*e[0], 1 - 1./3*e[0], -1 - 1./3*e[0]],
		 [-1 - 1./3*e[0], 2./3*e[0], 1 - 1./3*e[0]],
		 [1 - 1./3*e[0], -1 - 1./3*e[0], 2./3*e[0]] ])
#reward matrix for agent 2
rewardY = array( [ [2./3*e[1], 1 - 1./3*e[1], -1 - 1./3*e[1]],
		   [-1 - 1./3*e[1], 2./3*e[1], 1 - 1./3*e[1]],
		   [1 - 1./3*e[1], -1 - 1./3*e[1], 2./3*e[1]] ])

#Initialize empty lists to store plot points
#X and Y coordinates for agent X
xValsX = []
xValsY = []
#X and Y coordinates for agent Y
yValsX = []
yValsY = []

#Plot the triangular state space for each agent
graph.plot([-.5,0.,.5,-.5],[0.,.5*sqrt(3.),.0,.0])
graph.plot([-.5+1.,0.+1.,.5+1.,-.5+1.],[0.,.5*sqrt(3.),.0,.0])
canvas.show()

master.mainloop()