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

print '____________INSTRUCTIONS_______________'
print 'Use q,w,a,s,z & x keys to move arm'
print 'Press ESCAPE to exit'
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, 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.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 PointActualDAL(self, x, y): return [x - DISP_WIDTH/2, DISP_HEIGHT - y,0.0,1.0]

  # update the joint angels with the new values for joint indices in the list
  # e.g. updateArm([new_angle2, new_angle3], [1,2])
	def updateArm(self, angles, indices):
		for i in range(0, len(angles)):
			self.angles[indices[i]] = angles[i]

		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))]  

	def draw(self):
		self.screen.fill ((255, 255, 255)) #clean the screen first
		# draw the whole arm
		old_pts = [self.PointActualDAL(point[0], point[1]) 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]
		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:
			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

if __name__ == "__main__":

	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
	colors = [(0,0,255), (0,255,80), (255,0,0)] #the color for each link

	arm = ArmPlanning(length, np.radians(init_angles), colors)
	arm.addObstacle( (100,25,60), (-25, 75), int(12.5) )
	arm.draw()

	# Followings are testing
	ang1 = arm.angles[0]
	ang2 = arm.angles[1]
	ang3 = arm.angles[2]

  
	while not arm.checkCollision() and not arm.hit_ground:

		for event in pygame.event.get():
			if event.type == KEYUP and event.key == K_ESCAPE:
				pygame.quit()
		
		time.sleep(0.02)

		pressed = pygame.key.get_pressed()

		if pressed[K_q]:
			ang1 += 0.01
		if pressed[K_w]:
			ang1 -= 0.01
		if pressed[K_a]:
			ang2 += 0.01
		if pressed[K_s]:
			ang2 -= 0.01
		if pressed[K_z]:
			ang3 += 0.01
		if pressed[K_x]:
			ang3 -= 0.01

		#ang1 += 0.001
		#ang2 += 0.004
		#ang3 += 0.001
		arm.updateArm([ang1, ang2, ang3], [0,1,2])
		arm.draw()