import pygame
import math
import time
import numpy as np
import copy
from random import uniform
from PathGraph import *
from Queue import PriorityQueue
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 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

# Constants
# The size of the display window
DISP_WIDTH = 400
DISP_HEIGHT = 300
GRID_SIZE = 10
PLOT_CSPACE = True

# Robotic Arm Motion Planning Class
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.angles_ik = [0,0,0]
    self.angles_prev = [0,0,0]
    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.points_ik = [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.lines_ik = [pygame.rect.Rect(0,0,0,0)] * self.dof
    self.obstacle_rects = []
    self.hit_ground = False
    self.hit_ground_ik = False
    self.goal = None

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

  # Convert a point, (x,y) Tuple, to actual coordinate
  def PointActual(self, x, y): return (x - DISP_WIDTH/2, DISP_HEIGHT - y)

  # 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]
  
  # Draw all the components on the Canvas
  # Note: You need to draw things before checking Collision
  def draw(self):
    self.screen.fill ((255, 255, 255)) #clean the screen first
    self.hit_ground = False
    # draw the whole arm
    for i in range(0, self.dof):
      real_point = self.PointActual(self.points[i][0], self.points[i][1])
      # the angle for each joint is the sum of all the previous joints' angles
      self.points[i+1] = self.PointDisplay(real_point[0]+self.length[i]*math.cos(math.radians(sum(self.angles[:i+1]))),
                                           real_point[1]+self.length[i]*math.sin(math.radians(sum(self.angles[:i+1]))))

      self.lines[i] = pygame.draw.line(self.screen, self.colors[i], self.points[i], self.points[i+1], 5)
      #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)
    # draw the goal
    if self.goal != None:
      pygame.draw.rect(self.screen, [0,0,0], (self.goal[0]-5,self.goal[1]-5,10,10))
    pygame.display.update()


  # IK solver uses this to draw virtual arms and check collision
  # Note: You need to draw things before checking CollisionIK
  def draw_virtual(self):
    self.screen.fill ((255, 255, 255)) #clean the screen first
    self.hit_ground_ik = False
    # draw the whole arm
    for i in range(0, self.dof):
      real_point = self.PointActual(self.points[i][0], self.points[i][1])
      #still want to draw the real arms
      self.points[i+1] = self.PointDisplay(real_point[0]+self.length[i]*math.cos(math.radians(sum(self.angles[:i+1]))),
                                           real_point[1]+self.length[i]*math.sin(math.radians(sum(self.angles[:i+1]))))
      self.lines[i] = pygame.draw.line(self.screen, self.colors[i], self.points[i], self.points[i+1], 5)
      #virtual arms for IK solver
      self.points_ik[i+1] = self.PointDisplay(real_point[0]+self.length[i]*math.cos(math.radians(sum(self.angles_ik[:i+1]))),
                                           real_point[1]+self.length[i]*math.sin(math.radians(sum(self.angles_ik[:i+1]))))
      self.lines_ik[i] = pygame.draw.line(self.screen, (255,255,255), self.points_ik[i], self.points_ik[i+1], 5)
      #update the collision state with ground for the virtual arms
      if self.points_ik[i+1][1] > DISP_HEIGHT:
        self.hit_ground_ik = 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)
    # draw the goal
    if self.goal != None:
      pygame.draw.rect(self.screen, [0,0,0], (self.goal[0]-5,self.goal[1]-5,10,10))
    pygame.display.update()

  # Add an obstacle
  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))

  # Update the goal
  def addGoal(self, loc):
    self.goal = self.PointDisplay(loc[0],loc[1])

  # check all links for collision with ground or obstacles
  # Return True if any is in collision
  # Return False if collision free
  def checkCollision(self):
    if self.hit_ground:
      return True
    for l in self.lines:
      if (l.collidelist(self.obstacle_rects) != -1):
        return True
    return False

  # (For IK solver)
  # check all virtual links for collision with ground or obstacles
  # Return True if any is in collision
  # Return False if collision free
  def checkCollisionIK(self):
    if self.hit_ground_ik:
      return True
    for l in self.lines_ik:
      if (l.collidelist(self.obstacle_rects) != -1):
        return True
    return False

  # 2D Transition Matrix
  # pass ik=True to compute for virtual arms
  def A(self, i, ik=False): 
    if not ik:
      return np.array([[math.cos(math.radians(self.angles[i])),
                      -math.sin(math.radians(self.angles[i])), 
                      self.length[i]*math.cos(math.radians(self.angles[i]))] ,
                      [math.sin(math.radians(self.angles[i])), 
                      math.cos(math.radians(self.angles[i])), 
                      self.length[i]*math.sin(math.radians(self.angles[i]))] ,
                      [0, 0, 1]
                    ])
    else:
      return np.array([[math.cos(math.radians(self.angles_ik[i])),
                      -math.sin(math.radians(self.angles_ik[i])), 
                      self.length[i]*math.cos(math.radians(self.angles_ik[i]))] ,
                      [math.sin(math.radians(self.angles_ik[i])), 
                      math.cos(math.radians(self.angles_ik[i])), 
                      self.length[i]*math.sin(math.radians(self.angles_ik[i]))] ,
                      [0, 0, 1]
                    ])

  # Homogeneous Transformation from frame1 to frame2
  # pass ik=True to compute for virtual arms
  def T(self, f1, f2, ik=False):
    if f2 > self.dof or f1 < 0:
      raise ValueError("exceed the DOF when computing T!")
    t = np.eye(3)
    for k in range(0, f2-f1):
      if not ik:
        t = np.dot(t, self.A(k))
      else:
        t = np.dot(t, self.A(k,True))
    return t

  # Compute the current Jacobian Matrix
  # Jacobian Matrix Dimension: 6xN
  # N is the num of joints (DOFs)
  # we have a 6x3 matrix for this DOF-3 arm
  # pass ik=True to compute for virtual arms
  def J(self, ik=False):
    O0 = np.array([0,0,0])
    O1 = self.T(0,1,ik)[:,2].reshape(3)
    O2 = self.T(0,2,ik)[:,2].reshape(3)
    O3 = self.T(0,3,ik)[:,2].reshape(3)
    O1[2] = 0
    O2[2] = 0
    O3[2] = 0
    z = np.array([0,0,1])
    J1 = np.cross(z, O3-O0).reshape(3,1)
    J2 = np.cross(z, O3-O1).reshape(3,1)
    J3 = np.cross(z, O3-O2).reshape(3,1)
    J_up = np.hstack((J1, J2, J3))
    z = z.reshape(3,1)
    J_down = np.hstack((z,z,z))
    J = np.vstack((J_up, J_down))
    return J

  # Iterative Jacobian-Transpose IK Solver
  # input: target position
  # output: joint configuration
  # It first seeds using the current arm pose
  # If the solution is in collision, re-seed randomly
  # until collision-free solution is found, or reach MAX_Iter
  def Jtrans_IK(self, target):
    
    # IK solver parameters
    MAX_Iter = 20
    Error = 1
    MAX_STEPS = 100 #second

    self.angles_ik = copy.copy(self.angles) #seed with current angles
    count = 0
    while count < MAX_Iter:
      error = float('Inf')
      steps = 0
      while error > Error and steps < MAX_STEPS:
        J = self.J(True)[0:3,:]
        s = self.T(0,3,True)[:,2].reshape(3)
        s[2] = 0
        e = (target - s).reshape(3,1)
        alpha = 0.001

        # (3x3) x (3x1) = 3x1
        dTheta = alpha* np.dot(J.T, e)
        self.angles_ik[0] += dTheta[0][0]
        self.angles_ik[1] += dTheta[1][0]
        self.angles_ik[2] += dTheta[2][0]

        self.draw_virtual() #draw the virtual arm to step
        error = sum(sum(e**2))
        steps += 1
        
      #check if the solution is valid (collision free)
      #if self.checkCollisionIK() or self.hit_ground_ik:
      if self.checkCollisionIK():
        print 'Bad IK Solution'
        #re-seed with randomized angles
        self.angles_ik = [uniform(0,180), uniform(-180,180), uniform(-180,180)]
        print "re-seed the IK solver with ", self.angles_ik
        count += 1
        print 'Iter #', count, '\n'
      else:
        print "Found Good IK Solution"
        print "solution: ", self.angles_ik
        # check range
        for i in range(1,3):
          if self.angles_ik[i] > 180:
            self.angles_ik[i] = 360 - self.angles_ik[i]
          elif self.angles_ik[i] < -180:
            self.angles_ik[i] = 360 + self.angles_ik[i]
        print "bounded solution: ", self.angles_ik
        return copy.copy(self.angles_ik)

    # can not find solution
    print "Can not find valid IK solution"
    return None

  # Left-Click to set a goal for IK solver (any-time)
  # then visualize the steps through the goal
  # input: target position
  # output: joint configuration
  def Jtrans_IK_Mouse(self):
    #default target
    target = np.array([0, 250, 0])
    E = 1
    running = 1
    while running:
      # check left-click, update target
      event = pygame.event.poll()
      if event.type == pygame.KEYDOWN:
        running = 0
      elif event.type == pygame.MOUSEBUTTONDOWN:        
        t = self.PointActual(event.pos[0],event.pos[1])
        target[0] = t[0]
        target[1] = t[1]
        print "Goal: ", target

      J = self.J()[0:3,:]
      s = self.T(0,3)[:,2].reshape(3)
      s[2] = 0
      # Error vector (3x1)
      e = (target - s).reshape(3,1)
      alpha = 0.001

      # (3x3) x (3x1) = 3x1
      dTheta = alpha* np.dot(J.T, e)
      self.angles[0] += dTheta[0][0]
      self.angles[1] += dTheta[1][0]
      self.angles[2] += dTheta[2][0]
      self.draw()

      # check collision after drawing
      if arm.checkCollision():
        print "!! Collision -- Obstacle !!"
      elif arm.hit_ground:
        print "!! Collision -- Ground !!"

  # Create a graph object for the Configuration Space
  def createGraphConfig(self):

    nodeList = []
    nodeDict = {}
    x_free = y_free = z_free = np.array([])
    x_occ = y_occ = z_occ = np.array([])

    angles1 = range(0,181,GRID_SIZE)
    angles2 = range(-180,181,GRID_SIZE)
    angles3 = range(-180,181,GRID_SIZE)

    for k in range(len(angles3)):
      for j in range(len(angles2)):
        for i in range(len(angles1)):
          # try to move there and see if collision
          angles = [0+GRID_SIZE*i, -180+GRID_SIZE*j, -180+GRID_SIZE*k]
          self.updateArm(angles, [0,1,2])
          self.draw()

          if self.checkCollision():
            collide = True
            if PLOT_CSPACE:
              x_occ = np.append(x_occ, 0+GRID_SIZE*i)
              y_occ = np.append(y_occ, -180+GRID_SIZE*j)
              z_occ = np.append(z_occ, -180+GRID_SIZE*k)
          else:
            collide = False
            if PLOT_CSPACE:
              x_free = np.append(x_free, 0+GRID_SIZE*i)
              y_free = np.append(y_free, -180+GRID_SIZE*j)
              z_free = np.append(z_free, -180+GRID_SIZE*k)

          nid = len(nodeList)
          coords = (i,j,k)
          nodeDict[coords] = nid
          adjs = []
          
          #look at neighbors already seen
          # 13 possibilities:
          if (i-1,j-1,k) in nodeDict :
            adjs.append( (nodeDict[ (i-1,j-1,k) ], 1) ) #Add ajacency for this node
            nodeList[ nodeDict[ (i-1,j-1,k) ] ].adjs.append( (nid, 1) ) #Add adjacency for other node
          if (i-1,j,k) in nodeDict :
            adjs.append( (nodeDict[ (i-1,j,k) ], 1.0) ) #Add ajacency for this node
            nodeList[ nodeDict[ (i-1,j,k) ] ].adjs.append( (nid, 1.0) ) #Add adjacency for other node
          if (i,j-1,k) in nodeDict :
            adjs.append( (nodeDict[ (i,j-1,k) ], 1.0) ) #Add ajacency for this node
            nodeList[ nodeDict[ (i,j-1,k) ] ].adjs.append( (nid, 1.0) ) #Add adjacency for other node
          if (i+1,j-1,k) in nodeDict :
            adjs.append( (nodeDict[ (i+1,j-1,k) ], 1) ) #Add ajacency for this node
            nodeList[ nodeDict[ (i+1,j-1,k) ] ].adjs.append( (nid, 1) ) #Add adjacency for other node

          if (i-1,j-1,k-1) in nodeDict :
            adjs.append( (nodeDict[ (i-1,j-1,k-1) ], 1) ) #Add ajacency for this node
            nodeList[ nodeDict[ (i-1,j-1,k-1) ] ].adjs.append( (nid, 1) ) #Add adjacency for other node
          if (i-1,j,k-1) in nodeDict :
            adjs.append( (nodeDict[ (i-1,j,k-1) ], 1) ) #Add ajacency for this node
            nodeList[ nodeDict[ (i-1,j,k-1) ] ].adjs.append( (nid, 1) ) #Add adjacency for other node
          if (i-1,j+1,k-1) in nodeDict :
            adjs.append( (nodeDict[ (i-1,j+1,k-1) ], 1) ) #Add ajacency for this node
            nodeList[ nodeDict[ (i-1,j+1,k-1) ] ].adjs.append( (nid, 1) ) #Add adjacency for other node
          if (i,j-1,k-1) in nodeDict :
            adjs.append( (nodeDict[ (i,j-1,k-1) ], 1) ) #Add ajacency for this node
            nodeList[ nodeDict[ (i,j-1,k-1) ] ].adjs.append( (nid, 1) ) #Add adjacency for other node
          if (i,j,k-1) in nodeDict : 
            adjs.append( (nodeDict[ (i,j,k-1) ], 1.0) ) #Add ajacency for this node
            nodeList[ nodeDict[ (i,j,k-1) ] ].adjs.append( (nid, 1.0) ) #Add adjacency for other node
          if (i,j+1,k-1) in nodeDict :
            adjs.append( (nodeDict[ (i,j+1,k-1) ], 1) ) #Add ajacency for this node
            nodeList[ nodeDict[ (i,j+1,k-1) ] ].adjs.append( (nid, 1) ) #Add adjacency for other node
          if (i+1,j-1,k-1) in nodeDict :
            adjs.append( (nodeDict[ (i+1,j-1,k-1) ], 1) ) #Add ajacency for this node
            nodeList[ nodeDict[ (i+1,j-1,k-1) ] ].adjs.append( (nid, 1) ) #Add adjacency for other node
          if (i+1,j,k-1) in nodeDict :
            adjs.append( (nodeDict[ (i+1,j,k-1) ], 1) ) #Add ajacency for this node
            nodeList[ nodeDict[ (i+1,j,k-1) ] ].adjs.append( (nid, 1) ) #Add adjacency for other node
          if (i+1,j+1,k-1) in nodeDict :
            adjs.append( (nodeDict[ (i+1,j+1,k-1) ], 1) ) #Add ajacency for this node
            nodeList[ nodeDict[ (i+1,j+1,k-1) ] ].adjs.append( (nid, 1) ) #Add adjacency for other node

          # finally add to the nodeList
          nodeList.append(PathGraph.PathGraphNode(nid, coords, adjs, []))

          # check collision in C-Space
          if self.checkCollision():
            nodeList[nid].setState(True, 4)
            if PLOT_CSPACE:
              self.x_occ = np.append(x_occ, 0+GRID_SIZE*i)
              self.y_occ = np.append(y_occ, -180+GRID_SIZE*j)
              self.z_occ = np.append(z_occ, -180+GRID_SIZE*k)
          else:
            nodeList[nid].setState(False, 4)
            if PLOT_CSPACE:
              self.x_free = np.append(x_free, 0+GRID_SIZE*i)
              self.y_free = np.append(y_free, -180+GRID_SIZE*j)
              self.z_free = np.append(z_free, -180+GRID_SIZE*k)

    print "size of the graph: ", len(nodeList)
    graph = PathGraph(3, nodeList, nodeDict)
    self.graph = graph

    return graph

  # Plot the C-Space and Mark the current configuration
  def plotCSpace(self):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(self.x_free, self.y_free, self.z_free, zdir='z', s=40, c='b', marker='s')
    
    nid_prev = self.findNid(self.angles_prev)
    i,j,k = self.graph.getNode(nid_prev).getCoordinates()
    ax.scatter(np.array([0+GRID_SIZE*i]), np.array([-180+GRID_SIZE*j]), np.array([-180+GRID_SIZE*k]), 
               zdir='z', s=100, c='y', marker='s')

    nid_curr = self.findNid(self.angles)
    i,j,k = self.graph.getNode(nid_curr).getCoordinates()
    ax.scatter(np.array([0+GRID_SIZE*i]), np.array([-180+GRID_SIZE*j]), np.array([-180+GRID_SIZE*k]), 
               zdir='z', s=100, c='g', marker='o')

    x = y = z = np.array([])
    for nid in self.path:
      i,j,k = self.graph.getNode(nid).getCoordinates()
      x = np.append(x, 0+GRID_SIZE*i)
      y = np.append(y, -180+GRID_SIZE*j)
      z = np.append(z, -180+GRID_SIZE*k)
    ax.plot(x, y, z, label='parametric curve', c='r', linewidth=4)

    ax.set_title('C-Space Grids')
    ax.set_xlabel('Joint 1 (deg)')
    ax.set_ylabel('Joint 2 (deg)')
    ax.set_zlabel('Joint 3 (deg)')
    plt.show()

  # input: joint angles
  # output: nid in the grapg
  def findNid(self, angles):
    idx = (np.array(angles) - np.array([0, -180, -180]) ) / GRID_SIZE
    for i in range(0, len(idx)):
      idx[i] = int(idx[i])
    i = idx[0]
    j = idx[1]
    k = idx[2]
    p1 = 180. / GRID_SIZE + 1
    p2 = 360. / GRID_SIZE + 1
    nid = int((p2*p1)*k + (p1)*j + i)
    return nid


  def RRT(self, target_angles):
    self.graph


  # Returns the euclidean distance between two node id's in the given graph
  def euclidDist(self, a, b, graph):
    sum = 0
    for i in range(graph.coords):
      sum += (graph.getNode(a).coords[i] - graph.getNode(b).coords[i]) ** 2
    return float(sqrt(sum))


  # Keep track of "toVisit" nodes                
  # Node state: [distance, previousNodeId, visited, g+h]
  def AStar(self, start, goal):

    toVisit = [] # a list to keep track of potential next nodes

    # some initialization
    for i in range(len(self.graph.nodeList)):
      self.graph.nodeList[i].state[0] = float("Inf") #distance
      self.graph.nodeList[i].state[1] = -1 #previousNodeId
      self.graph.nodeList[i].state[2] = False #have been added to toVisit or not?
      self.graph.nodeList[i].state[3] = float("Inf")

    # set the start node's distance to zero
    self.graph.nodeList[start].state[0] = 0 #distance is zero
    self.graph.nodeList[start].state[3] = self.euclidDist(start, goal, self.graph)
    self.graph.nodeList[start].state[2] = True
    toVisit.append(start)

    # main loop
    while (1):

      #find the nodes with smallest distance
      minNode, minNodeId, minDist, idx = self.findSmallestDist(toVisit)
      # The goal has been reached
      if minNodeId == goal:
        break

      # remove this node from toVisit set
      print 'len(toVisit)=', len(toVisit)
      print 'idx = ', idx
      del toVisit[idx]

      #compute the distance for all neighbors
      for adj in minNode.getAdjacencies() :

        adj_nid = adj[0]
        print 'adj_nid= ', adj_nid
        adj_dist = adj[1]
        adj_node = self.graph.getNode(adj_nid)
        print 'Collide? ', adj_node.state[4]
        if adj_node.state[2] == False and adj_node.state[4] == False:
          toVisit.append(adj_nid) #want to keep track of this nodes
          adj_node.setState(True, 2)
        # Improved distance
        h = self.euclidDist(adj_nid, goal, self.graph)
        g = minNode.state[0] + adj_dist

        if (g+h < adj_node.state[3]):
          adj_node.setState(g, 0) #update the distance
          adj_node.setState(minNodeId, 1) #update previousNodeId
          #adj_node.setState(max(g+h, minNode.state[3]), 3) #update the g+h
          adj_node.setState(g+h, 3) #update the g+h

    # Form the path    
    self.formPath(start, goal)
    return self.path

  #Find the node with smallest distance in the unvisited set
  def findSmallestDist(self, toVisit):

    minNodeId = 0
    minDist = float("Inf")
    idx = 0
    count = 0
    for i in toVisit:
      if self.graph.nodeList[i].state[3] < minDist:
        idx = count
        minDist = self.graph.nodeList[i].state[3]
        minNodeId = i
      count += 1
    minNode = self.graph.nodeList[minNodeId]
    return minNode, minNodeId, minDist, idx

  # Reverse the Nodes and get the path 
  def formPath(self, start, goal):
    self.path = []
    self.path.append(goal)
    current_nid = goal
    while current_nid != start:
      self.path.insert(0, self.graph.nodeList[current_nid].state[1])
      print "current id=" ,self.graph.nodeList[current_nid].state[1]
      current_nid = self.graph.nodeList[current_nid].state[1]
    print self.path

  # A-Star pathfinding algorithm.
  # A greedy algorithm that uses an admissable heuristic, so it finds a guaranteed shortest path.
  def AStar2(self, start, goal):

    self.visited = []
    if self.graph != None:
      for i in range(self.graph.getNodeCount()):
        self.visited.append(None)

    toVisit = PriorityQueue()

    # Start of the algorithm
    self.visited[start] = (0, None)
    toVisit.put( (0 + self.euclidDist(start, goal, self.graph), start, None) )

    while not toVisit.empty(): # Start of the main loop

      nextTuple = toVisit.get() #Get the next node

      next = nextTuple[1]
      nextDist = self.visited[next][0] #Total distance travelled to get to here.
      #Use this instead of the toVisit, as that is estimated dist. to goal 

      if next == goal: #Success!
        lastPathNode = goal # Construct the path list
        self.pathCost = self.visited[goal][0]  # 0 denotes cost to get to this node
        self.path = []
        while self.visited[lastPathNode][1] != None:
          self.path.append(lastPathNode)
          lastPathNode = self.visited[lastPathNode][1] # 1 denotes the node we just came from
        self.path.append(start)
        self.path.reverse()
        print 'The Planned Path:', self.path
        return self.path

      for adjacency in self.graph.getAdjacencies(next): # For each neightbor, see if have a path to
        nextNbr, cost = adjacency[0], adjacency[1]
        totalTravel = nextDist + cost
        if not self.graph.nodeList[nextNbr].state[4]: #if not collision
          if self.visited[nextNbr] == None: # The neighbor is unvisited
            self.visited[nextNbr] = ( totalTravel, next ) # Mark the node as visited
            goalDist = self.euclidDist(nextNbr, goal, self.graph) # Estimated distance (euclidean)
            toVisit.put( (totalTravel + goalDist, nextNbr, next) ) # Add node to priority queue
          elif totalTravel < self.visited[nextNbr][0]: # The neighbor has been visited, but that is shorter
                                                        # if we've found a shorter route to nextNbr
            self.visited[nextNbr] = (totalTravel, next) # then replace old route to nextNbr with new
            goalDist = self.euclidDist(nextNbr, goal, self.graph)
            toVisit.put( (totalTravel + goalDist, nextNbr, next) )
    #else: Node has been visited, but current path is equal or longer
    self.path = None
    print 'No Collision Free Path is found, Please try again'
    return None


if __name__ == "__main__":

  # define the arm parameters
  length = [100, 100, 50] #the length of links
  init_angles = [90, 0, 0] #joint angles
  colors = [(0,0,255), (0,255,80), (255,0,0)] #the color for each link
  arm = ArmPlanning(length, init_angles, colors)

  # put obstacles
  arm.addObstacle( (50,0,100), (-110, 180), int(30) )
  arm.addObstacle( (100,25,60), (90, 200), int(30) )
  #arm.addObstacle( (50,0,100), (140, 120), int(25) )

  # explore and create graph in C-Space
  arm.createGraphConfig()

  # draw initial pose
  init_angles = [90, 0, 0] #joint angles
  start_id = arm.findNid(init_angles)
  arm.updateArm([init_angles[0],init_angles[1],init_angles[2]], [0,1,2])
  arm.draw()

  # Main Loop
  running = 1
  while running:
    # check left-click, update goal
    event = pygame.event.poll()
    if event.type == pygame.KEYDOWN:
      if event.key == pygame.K_ESCAPE: #press ESC to leave
        running = 0
      elif event.key == pygame.K_d:
        arm.plotCSpace()
    if event.type == pygame.MOUSEBUTTONDOWN:        
      t = arm.PointActual(event.pos[0],event.pos[1])
      target = np.array([t[0], t[1], 0])
      print "New Goal: ", target
      arm.addGoal([t[0],t[1]])

      # Find IK solution
      print '\n@@@@@@@@@@@@@ FINDING IK SOLUTION...... @@@@@@@@@@@@@@\n'
      target_angles = arm.Jtrans_IK(target)
      if target_angles == None:
        print '==== Try Another Goal ! ===='      
      else:
        target_id = arm.findNid(target_angles)
        print 'target_id= ', target_id

        # current nid
        start_id = arm.findNid(arm.angles)

        # planning in C-Space
        print '\n@@@@@@@@@@@@@@ PLANNING...... @@@@@@@@@@@@@\n'
        path = arm.AStar2(start_id, target_id)
        if path == None:
          print '===== Try Again ! ======'

        else:
          arm.angles_prev = copy.copy(arm.angles)
          # move the arm according to the planning results
          for nid in path:
            time.sleep(0.05) #delay for doing animation
            i,j,k = arm.graph.getNode(nid).getCoordinates()
            angles = [0+GRID_SIZE*i, -180+GRID_SIZE*j, -180+GRID_SIZE*k]
            arm.updateArm([angles[0],angles[1],angles[2]], [0,1,2])
            arm.draw()

          # finished
          print "Goal achieved!"
          print 'final angle: ', arm.angles, ' , nid=', arm.findNid(arm.angles)