import urx, math, math3d, time


class UR5Exception(Exception):
	pass


class DemoController(object):
	"""Implements control of the UR5 in cylinder coordinates (corresponding to attached table).
	   
	-Checks bounds of the arc section
	-Methods for placing blocks
	"""

	def __init__(self, config):
		"""Initializes the controller with parameters specified in config"""
		self.debug = config.DEBUG
		try:
			self.robot = urx.Robot(config.UR5_IP, useRTInterface=config.USE_FORCE_MONITOR) #TODO or hostname
		except Exception, e:
			raise UR5Exception('Robot initiation failed: %s' % str(e))

		# Control params
		self.vel = config.VELOCITY
		self.acc = config.ACCELERATION
		self.dec = config.DECELERATION

		# Cylinder params
		self.quadrant = config.TABLE_QUADRANT
		self.phi = self.quadrant*math.pi/2
		self.cyl_offset = config.TABLE_ORIGO_OFFSET
		self.cyl_ox = self.cyl_offset*math.cos(self.phi)
		self.cyl_oy = self.cyl_offset*math.sin(self.phi)

		self.r_min = config.R_MIN
		self.r_max = config.R_MAX
		self.theta_min = config.THETA_MIN
		self.theta_max = config.THETA_MAX
		self.z_min = config.Z_MIN
		self.z_max = config.Z_MAX

		# Block params
		self.block_dim = config.BLOCK_DIM
		self.r_lvl0 = config.R_LVL0
		self.z_lvl0 = config.Z_LVL0
		self.r_bmargin = config.R_BLOCKMOVE_MARGIN
		self.theta_bmargin = config.THETA_BLOCKMOVE_MARGIN
		self.z_bmargin = config.Z_BLOCKMOVE_MARGIN
		self.r_boffset = config.R_BLOCKMOVE_OFFSET
		self.theta_boffset = config.THETA_BLOCKMOVE_OFFSET
		self.z_boffset = config.Z_BLOCKMOVE_OFFSET
		
		# Home quadrant joint config
		self.j_home = config.QUADRANT_JOINT_CONFIG
	

	def calibrate_csys(self):
		"""Calibrate the reference cylinder coordinate system."""
		# Move to reference coordinate system base
		self.set_pose_home()
		# Set reference coordinate system parameters
		csys = math3d.Transform()
		csys.orient.rotate_zb(self.phi)
		self.robot.set_csys("csys", csys)
		self.trans_base = self.robot.get_transform()
		self.cyl_ox = self.cyl_offset
		self.cyl_oy = 0


	def set_freedrive(self, mode):
		"""Set UR5 to freedrive mode."""
		self.robot.set_freedrive(mode)


	def set_pose_home(self):
		"""Move to "home" configuration. The resulting pose has correct orientation."""
		try:
			self.robot.movej(self.j_home, acc=0.1, vel=0.5, radius=0, wait=True, relative=False)
			pose = self.robot.getl()
			pose[0] = self.r_min*math.cos(self.phi) + self.cyl_ox
			pose[1] = self.r_min*math.sin(self.phi) + self.cyl_oy
			pose[2] = self.z_min
			self.robot.movel(pose, acc=0.1, vel=0.2, radius=0, wait=True, relative=False)
			self.current_cyl = (self.r_min, 0, self.z_min)
		except Exception, e:
			raise UR5Exception('Move to home configuration failed: %s' % str(e))


	def cylinder2pose(self, r, theta, z):
		"""Translate table cylinder coordinates to robot cartesian pose."""
		trans = math3d.Transform(self.trans_base) #deep copy
		trans.orient.rotate_zb(theta)
		trans.pos = math3d.Vector(r*math.cos(theta) + self.cyl_ox,
								  r*math.sin(theta) + self.cyl_oy,
								  z)
		return trans.pose_vector.tolist()


	def blocklvl2pose(self, r_lvl, theta, z_lvl):
		return self.cylinder2pose(self.r_lvl0 - r_lvl*self.block_dim, theta, self.z_lvl0 + z_lvl*self.block_dim)


	def movel(self, r, theta, z):
		"""Linear move."""
		pose = self.cylinder2pose(r, theta, z)
		self.robot.movel(pose, acc=self.acc, vel=self.vel, radius=0, wait=True, relative=False)
		self.current_cyl = (r, theta, z)


	def movec(self, r, theta, z):
		"""Circular move - uses URScript built-in function.

		Unfortunately, this doesn't work very well: The tool orientation at the target pose 
		is not right. URScript specifies that orientation in pose_via is not considered, however,
		it SHOULD consider the orientation of the end pose.
		"""
		curr_theta = self.current_cyl[1]
		pose_via = self.cylinder2pose(r, curr_theta+(theta-curr_theta)/2, z)
		pose = self.cylinder2pose(r, theta, z)
		self.robot.movec(pose_via, pose, acc=self.acc, vel=self.vel, radius=0, wait=True)
		self.current_cyl = (r, theta, z)


	def movels(self, move_list, wait=True):
		"""Concatenate several movel commands and applies a blending radius.
		move_list = pose + [acc, vel, radius]
		"""
		header = "def myProg():\n"
		end = "end\n"
		template = "movel(p[{},{},{},{},{},{}], a={}, v={}, r={})\n"
		prog = header + ''.join([template.format(*m) for m in move_list]) + end
		self.robot.send_program(prog)
		if not wait:
			return None
		else:
			self.robot.wait_for_move()
			return self.robot.getl()


	def movec_hax(self, r, theta, z):
		"""movec is unreliable and does not finish in the correct pose. We work around this
		by linearizing the arc into segments of a predefined angular resolution, and use
		movepls with blend.

		IMPORTANT: blending radius have to be smaller than the angular resolution! If not,
		the robot will not finish the last move, because it is within the target
		"""
		move_list = []
		step = math.pi/41
		curr_r, curr_theta, curr_z = self.current_cyl
		dtheta = theta - curr_theta

		segments = max(int(round(abs(dtheta)/step)), 1)
		theta_incr = dtheta/segments
		r_incr = (r-curr_r)/segments
		z_incr = (z-curr_z)/segments
		for i in range(1, segments):
			pose = self.cylinder2pose(curr_r + i*r_incr, curr_theta + i*theta_incr, curr_z + i*z_incr)
			move_list.append(pose + [self.acc, self.vel, 0.03])
		move_list.append(self.cylinder2pose(r, theta, z) + [self.acc, self.vel, 0])

		if self.debug:
			print "move list for movec_hax"
			for p in move_list:
				print [('%.3f' % i) for i in p]

		self.movels(move_list, wait=True)
		self.current_cyl = (r, theta, z)


	def pick_block(self, r_lvl, theta, z_lvl, speed):
		"""This function concatenates 3 moves:

		1. Move the tool in front of the block.
		2. Move the tool into the hole.
		3. Lift the block.

		Here, we assume that the tool is nearby the block. Bringing the tool towards the block is
		not the responsibility of this function.

		r_lvl starts at 0, which correspond to the outer edge of the table.
		z_lvl starts at 0, which corrsepond to table level.
		"""
		r_target = self.r_lvl0 - r_lvl*self.block_dim
		z_target = self.z_lvl0 + z_lvl*self.block_dim
		p1 = self.cylinder2pose(r_target - self.r_bmargin, theta - self.theta_bmargin, z_target - self.z_bmargin)
		p2 = self.cylinder2pose(r_target, theta, z_target)
		p3 = self.cylinder2pose(r_target + self.r_boffset, theta + self.theta_boffset, z_target + self.z_boffset)
		move_list = [
			p1 + [self.acc, speed, 0.005],
			p2 + [self.acc, speed, 0],
			p3 + [self.acc, speed, 0.01],
		]
		return move_list

	def place_block(self, r_lvl, theta, z_lvl, speed):
		"""Reverse move of pick_block."""
		moves = self.pick_block(r_lvl, theta, z_lvl, speed)
		moves.reverse()
		return moves


	def construct_moves(self, task_list, speed):
		"""Construct moves from a list of tasks"""
		pass # wait with implementing this


	def is_looping(self, dt, threshold=0.001):
		"""Polls the robot for change in pose vector. Blocks for dt seconds."""
		pose0 = self.robot.getl()
		time.sleep(dt)
		pose1 = self.robot.getl()
		v = [pose1[i]-pose0[i] for i in range(len(pose0))]
		if self.debug:
			_sum = sum(map(abs, v))
			print "Velocity vector, dt=%.2f: %s. SUM %.3f" % (dt, str([('%.2f' % i) for i in v]), _sum)
		return sum(map(abs, v)) > threshold

	def cleanup(self):
		self.robot.cleanup()