#!/usr/bin/env python
# -*- coding: iso-8859-15 -*-

"""
roslaunch darknet_ros darknet_ros.launch image:=/camera/color/image_raw
"""

import sys
import rospy
import tf
import tf2_ros
import tf2_geometry_msgs
import image_geometry
import moveit_commander
import numpy as np
from math import *
from time import sleep
from sensor_msgs.msg import Image, CameraInfo
from darknet_ros_msgs.msg import BoundingBoxes
from geometry_msgs.msg import PoseStamped
from cv_bridge import CvBridge

class OpenManipGrabBottle:
	def __init__(self):
		rospy.init_node('open_manip_grab_bottle')
		moveit_commander.roscpp_initialize(sys.argv)
		self.sub_depthinfo = rospy.Subscriber("/camera/aligned_depth_to_color/camera_info", CameraInfo, self.depthinfo_received)
		self.tf_broadcaster = tf.broadcaster.TransformBroadcaster()
		self.depthimgmsg = None
		self.camera_model = None
		self.bridge = CvBridge()
		self.tfBuffer = tf2_ros.Buffer()
		self.listener = tf2_ros.TransformListener(self.tfBuffer)
		self.target_class = "bottle"

		self.gripper_pose_pub = rospy.Publisher('gripper_pose', PoseStamped)
		self.arm	= moveit_commander.MoveGroupCommander("arm")
		self.gripper	= moveit_commander.MoveGroupCommander("gripper")
		self.arm.set_goal_position_tolerance(0.01)
		self.arm.set_goal_orientation_tolerance(5*pi/180)
		self.home()
		sleep(10) # wait for image
		self.register_subscriber()
		rospy.spin()

	def register_subscriber(self):
		self.sub_obj_detection = rospy.Subscriber("/darknet_ros/bounding_boxes", BoundingBoxes, self.detected_objects)
		self.sub_depthimage = rospy.Subscriber("/camera/aligned_depth_to_color/image_raw", Image, self.depthimage_received)

	def home(self):
		self.gripper.set_named_target('close')
		self.gripper.go()
		#self.arm.set_named_target('home')
		self.arm.go([f * pi/180 for f in [90.0, -90, 77, 20]], wait=True)

	def detected_objects(self, msg):
		if self.depthimgmsg is None or self.camera_model is None:
			return
		box_bottles = filter(lambda box: box.Class == self.target_class and box.probability > 0.5, msg.bounding_boxes)
		if len(box_bottles) > 0:
			box = box_bottles[0] # just use the first one found for now
			box_x = int(box.xmin + (box.xmax - box.xmin)/2.0)
			box_y = int(box.ymin + (box.ymax - box.ymin)/2.0)
			depthframe = self.bridge.imgmsg_to_cv2(self.depthimgmsg)
			dist = depthframe[box_y][box_x]/1000.0 # assume size rgb = size depth
			if dist == 0:
				rospy.logerr("depth invalid")
				return

			self.sub_obj_detection.unregister()
			self.sub_depthimage.unregister()

			coords3d = [f * dist for f in self.camera_model.projectPixelTo3dRay([box_x, box_y])]
			point = PoseStamped()
			point.header.frame_id = self.depthimgmsg.header.frame_id
			point.header.stamp = rospy.Time.now()
			point.pose.position.x = coords3d[0]
			point.pose.position.y = coords3d[1]
			point.pose.position.z = coords3d[2]
			point = self.tfBuffer.transform(point, "open_manipulator_base", rospy.Duration(1))
			point.pose.orientation.x, point.pose.orientation.y, point.pose.orientation.z, point.pose.orientation.w = [0, 0, 0, 1] # bottle stands up
			print point
			point.pose.position.z-=0.02

			self.tf_broadcaster.sendTransform([point.pose.position.x, point.pose.position.y, point.pose.position.z], [point.pose.orientation.x, point.pose.orientation.y, point.pose.orientation.z, point.pose.orientation.w], rospy.Time.now(), self.target_class, "open_manipulator_base")

			if self.pick(point):
				rospy.signal_shutdown("done")
			else:
				self.home()
				self.register_subscriber()

	def depthimage_received(self, msg):
		self.depthimgmsg = msg

	def depthinfo_received(self, msg):
		self.camera_model = image_geometry.PinholeCameraModel()
		self.camera_model.fromCameraInfo(msg)
		self.sub_depthinfo.unregister()

	def pick(self, pose):
		pos2d = [pose.pose.position.x, pose.pose.position.y]
		dist = np.linalg.norm(pos2d)
		pos2d_unit = pos2d / dist # unit vector

		# orientation towards base
		q = tf.transformations.quaternion_from_euler(0, 0, acos(np.dot(pos2d_unit, [1, 0]))) # use the dot product to get the angle
		pose.pose.orientation.x, pose.pose.orientation.y, pose.pose.orientation.z, pose.pose.orientation.w = [q[0], q[1], q[2], q[3]]

		# goal 5cm before
		pose.pose.position.x = pos2d_unit[0] * (dist-0.05)
		pose.pose.position.y = pos2d_unit[1] * (dist-0.05)
		self.arm.set_pose_target(pose)
		if not self.arm.go():
			print "before position failed"
			return False

		self.gripper.set_named_target('open')
		self.gripper.go()

		# goal on target +3cm
		pose.pose.position.x = pos2d_unit[0] * (dist+0.03)
		pose.pose.position.y = pos2d_unit[1] * (dist+0.03)
		self.arm.set_pose_target(pose)
		if not self.arm.go():
			print "position on target failed"
			return False

		self.gripper.set_named_target('close')
		self.gripper.go()

		# height
		angles = [0.0, 0.0, 0.0, 0.0]
		while angles == [0.0, 0.0, 0.0, 0.0]: # invalid
			angles = self.arm.get_current_joint_values()
		print [f * 180/pi for f in angles]
		angles[1]-=20*pi/180
		if not self.arm.go(angles, wait=True):
			print "higher position failed"
			return True

		if not self.arm.go([f * pi/180 for f in [-90, -60, 0, 0]], wait=True):
			print "final position failed"
			return True

		return True

if __name__ == "__main__":
	OpenManipGrabBottle()