import pygame
from pygame.locals import *
import time
import random
import numpy as np

print '____________INSTRUCTIONS________________'
print 'Move mouse to change end effector goal'
print 'Press ESCAPE to quit'
print '________________________________________'

# the units are (cm or pixel), and (degree) 
# The coordinate of PyGame is:
#    ---> x
#    '
#    '
#    y
# the upperleft pixel is (0,0)
# Note: use function PointDisplay() to convert to PyGame coordinate
#       use function PointActual()  to convert to the actual coordinate

# The size of the display window
DISP_WIDTH = 500
DISP_HEIGHT = 300

#cumulative dot product for a list of numpy arrays
def CumulDot(MList, i):
	if i == 0:
		Cdot = MList[0]
	else:
		Cdot = MList[0]
		for j in range(1,i+1):
			 Cdot = np.dot(Cdot,MList[j])
	return Cdot 



class ArmPlanning():

	def __init__(self, length, angles, goal, colors):
		pygame.init()
		self.screen = pygame.display.set_mode((DISP_WIDTH, DISP_HEIGHT))
		self.screen.fill((255,255,255))

		self.length = length
		self.angles = angles #joint angles
		self.colors = colors #the color for each link
		self.dof = len(angles) #degree of freedom
		self.points = [self.PointDisplay([0, 0])]*(self.dof+1) # shift to the center
		self.goal = goal

		self.obstacles = []
		self.lines = [pygame.rect.Rect(0,0,0,0)] * self.dof
		self.obstacle_rects = []
		self.hit_ground = False

		c = np.cos(angles)
		s = np.sin(angles)
		self.A = [np.array([[c[i], -s[i], 0 , length[i] * c[i]], [s[i], c[i], 0 , length[i] * s[i]],[0,0,1,0],[0,0,0,1]]) for i in range(len(length))]

		self.T = [CumulDot(self.A,i) for i in range(len(self.A))]     

	# Convert a point, (x,y) Tuple, to PyGame coordinate
	def PointDisplay(self, point): return (DISP_WIDTH/2 + point[0], DISP_HEIGHT - point[1])

	#my version, which includes z and 1 for use with D-H convention and uses lists for dot product
	def PointTrue(self, (x, y)): return [x - DISP_WIDTH/2, DISP_HEIGHT - y,0.0,1.0]

  # update the joint angels with the new values 
	def fKinematics(self, angles):
		self.angles = angles

		c = np.cos(angles)
		s = np.sin(angles)
		self.A = [np.array([[c[i], -s[i], 0 , self.length[i] * c[i]], [s[i], c[i], 0 , self.length[i] * s[i]],[0,0,1,0],[0,0,0,1]]) for i in range(len(self.length))]
		self.T = [CumulDot(self.A,i) for i in range(len(self.A))]
		
		old_pts = [self.PointTrue(point) for point in self.points]
		new_pts = [np.dot(self.T[i], old_pts[0]).tolist() for i in range(len(self.T))]
		new_pts.insert(0, old_pts[0])

		self.points = [self.PointDisplay(point) for point in new_pts]

	def draw(self):
		self.screen.fill ((255, 255, 255)) #clean the screen first
		# draw the whole arm
		self.lines = [pygame.draw.line(self.screen, colors[i], self.points[i], self.points[i+1], 5) for i in range(0, self.dof)] 

		#check if it collides with ground
		if self.points[i+1][1] > DISP_HEIGHT:
#			print 'arm hit ground'
			self.hit_ground = True

		# draw the obstacles
		for i in range(0, len(self.obstacles)):
			self.obstacle_rects[i] = pygame.draw.circle(self.screen, self.obstacles[i][0], self.obstacles[i][1], self.obstacles[i][2], 0)

		pygame.display.update()
	
	def addObstacle(self, color, loc, radius):
		obs = (color, self.PointDisplay([loc[0],loc[1]]), radius)
		self.obstacles.append(obs)
		self.obstacle_rects.append(pygame.rect.Rect(0,0,0,0))

	# check each link if it collides with ground or obstacles
	def checkCollision(self):
		for l in self.lines:
			if (l.collidelist(self.obstacle_rects) != -1):
				print 'Hit the obstacle!'
				return True
		return False

	def getJacobian(self):
		ai = np.array([0,0,1]) #unit length rotation axis in WORLD space (always z-axis for 2D problem)
		e2 = self.PointTrue(self.points[-1][:3])
		e2 = np.asarray(e2[:3])
		r2 = self.PointTrue(self.points[1][:3])
		r2 = r2[:3]

		e = np.array(self.points[-1][:3]) #position of end effector in WORLD space
		dedt = [np.cross(ai, e2 - np.asarray(self.PointTrue(point[:3])[:3])) for point in self.points[0:self.dof]]

		J1 = np.array(dedt).transpose()
		J2 = np.array([ai]*self.dof).transpose()
		J=np.vstack((J1,J2))
		return J


if __name__ == "__main__":

	length = [100, 50, 50,0]
	#length = [100, 120, 50] #the length of links
	#init_angles = [10, 15, 5] #joint angles  colors = [(0,0,255), (0,255,80), (255,0,0)] #the color for each link
	init_angles = [10,10,15,0]
	#colors = [(0,0,255), (0,255,80), (255,0,0)] #the color for each link
	colors = [(0,0,255),(255,0,0),(0,255,0),(0,0,0)]
	goal = np.array([0,200,0])
	
	arm = ArmPlanning(length, np.radians(init_angles), goal, colors)
	arm.addObstacle( (100,25,60), goal.tolist(), int(12.5) )
	arm.draw()
	
	beta = 0.01
	tolerance = 5
	i=0

	while True:
		for event in pygame.event.get():
			if event.type == KEYUP and event.key == K_ESCAPE:
				pygame.quit()
			elif event.type == pygame.MOUSEMOTION:
				goal = np.array(list(arm.PointTrue((event.pos)))[:3])


		location = arm.PointTrue(arm.points[-1])
		distance = np.linalg.norm(goal-location[:3])
		#while i < 5: 
		while distance > tolerance:

			for event in pygame.event.get():
				if event.type == KEYUP and event.key == K_ESCAPE:
					pygame.quit()
				elif event.type == pygame.MOUSEMOTION:
					goal = np.array(list(arm.PointTrue((event.pos)))[:3])

			J = arm.getJacobian()
			
			L=[0.100,0.20,100,10.01,0.01,100.01]
			#JI = np.dot(J.transpose(), np.linalg.inv(np.dot(J,J.transpose())+L*np.identity(6)))
			JI = np.linalg.pinv(J)
			#JI=J.transpose()
			goal3D = np.hstack((goal,[0,0,1]))
			loc3D = np.hstack((location[:3],[0,0,1]))
			de = beta * (goal3D-loc3D)
			dt = np.dot(JI,de)
			theta = arm.angles + dt
			arm.fKinematics(theta)
			arm.draw()
			i += 1
			if i>500:
				print 'failed to converge'
				break
			location = arm.PointTrue(arm.points[-1])
			distance = np.linalg.norm(goal-np.asarray(location[:3])) 

			arm.draw()

	print 'made it to target'	
	print 'final location: ' + str(arm.PointTrue(arm.points[-1]))
	print 'goal: ' + str(goal)