import sys
import rospy
import moveit_commander
import moveit_msgs.msg
import geometry_msgs.msg

import include.io_control_lib as io
import include.yaml_data_handling as yaml_handle

# Debug mode, when True, show a lot of useful messages
DEBUG = rospy.get_param("robocheckers/debug")


class MoveGroup(object):
    def __init__(self):
        super(MoveGroup, self).__init__()

        # First initialize `moveit_commander`_ and a `rospy`_ node:
        moveit_commander.roscpp_initialize(sys.argv)
        rospy.init_node('robot_move')

        # Instantiate a `RobotCommander`_ object. Provides information such as the robot's
        # kinematic model and the robot's current joint states
        robot = moveit_commander.RobotCommander()

        # Instantiate a `PlanningSceneInterface`_ object.    This provides a remote interface
        # for getting, setting, and updating the robot's internal understanding of the
        # surrounding world:
        scene = moveit_commander.PlanningSceneInterface()

        # This interface can be used to plan and execute motions:
        group_name = "manipulator"
        move_group = moveit_commander.MoveGroupCommander(group_name)

        # Create a `DisplayTrajectory`_ ROS publisher which is used to display
        # trajectories in Rviz:
        display_trajectory_publisher = rospy.Publisher('/move_group/display_planned_path',
                                                       moveit_msgs.msg.DisplayTrajectory,
                                                       queue_size=20)

        # Creating the publisher for the io_control
        io_control = rospy.Publisher('/robocheckers/io_message', io.String, queue_size=10)

        # We can get the name of the reference frame for this robot:
        planning_frame = move_group.get_planning_frame()
        if DEBUG: print("============ Planning frame: %s" % planning_frame)

        # We can also print(the name of the end-effector link for this group:)
        eef_link = move_group.get_end_effector_link()
        if DEBUG: print("============ End effector link: %s" % eef_link)

        # We can get a list of all the groups in the robot:
        group_names = robot.get_group_names()
        if DEBUG: print("============ Available Planning Groups:", robot.get_group_names())

        # Sometimes for debugging it is useful to print(the entire state of the)
        # robot:
        if DEBUG:
            print("============ Printing robot state")
            print(robot.get_current_state())
            print("")

        # Misc variables
        self.robot = robot
        self.scene = scene
        self.move_group = move_group
        self.display_trajectory_publisher = display_trajectory_publisher
        self.planning_frame = planning_frame
        self.eef_link = eef_link
        self.group_names = group_names
        self.io_control = io_control

        # Declaring dictionary with static position values
        self.static_positions = {}

        # Inner state telling how many kings have already been taken
        self.king_count = 0


    # Function for adding virtual robot cage for collision checking
    def addCollisionObjects(self):
        # Defining size
        WIDTH = 0.65
        LENGTH = 1.15
        HEIGHT = 0.85
        WALL_WIDTH = 0.005

        # Waiting for scene and robot init unless objects wont appear
        rospy.sleep(2)

        # Defining coordinate system
        wall = geometry_msgs.msg.PoseStamped()
        wall.header.frame_id = self.robot.get_planning_frame()
        wall.pose.orientation.w = 1

        # Front wall
        wall.pose.position.x = LENGTH/2+0.21-WALL_WIDTH/2
        wall.pose.position.y = 0.
        wall.pose.position.z = HEIGHT/2
        self.scene.add_box("front_wall", wall, (WALL_WIDTH, WIDTH-2*WALL_WIDTH, HEIGHT))

        # Rear wall
        wall.pose.position.x = -LENGTH/2+0.21+WALL_WIDTH/2
        wall.pose.position.y = 0.
        wall.pose.position.z = HEIGHT/2
        self.scene.add_box("rear_wall", wall, (WALL_WIDTH, WIDTH-2*WALL_WIDTH, HEIGHT))

        # Left wall
        wall.pose.position.x = 0.21
        wall.pose.position.y = (WIDTH-WALL_WIDTH)/2
        wall.pose.position.z = HEIGHT/2
        self.scene.add_box("left_wall", wall, (LENGTH, WALL_WIDTH, HEIGHT))

        # Right wall
        wall.pose.position.x = 0.21
        wall.pose.position.y = -(WIDTH-WALL_WIDTH)/2
        wall.pose.position.z = HEIGHT/2
        self.scene.add_box("right_wall", wall, (LENGTH, WALL_WIDTH, HEIGHT))

        # Ceiling
        wall.pose.position.x = 0.21
        wall.pose.position.y = 0
        wall.pose.position.z = HEIGHT+WALL_WIDTH/2
        self.scene.add_box("ceiling", wall, (LENGTH, WIDTH, WALL_WIDTH))

        # Chessboard
        wall.pose.position.x = 0.36
        wall.pose.position.y = 0.
        wall.pose.position.z = 0.05
        self.scene.add_box("chessboard", wall, (0.4, 0.4, 0.02))


    # Function for moving robot to desired locations
    def moveAboveBoard(self, pos):
        pose_goal = geometry_msgs.msg.Pose()
        # Setting coordinates of piece
        pose_goal.position.x = pos[0]
        pose_goal.position.y = pos[1]
        pose_goal.position.z = 0.3

        # Setting orientation
        pose_goal.orientation.x = 1
        pose_goal.orientation.z = 0

        self.move_group.set_pose_target(pose_goal)

        plan = self.move_group.plan()

        display_trajectory = moveit_msgs.msg.DisplayTrajectory()
        display_trajectory.trajectory_start = self.robot.get_current_state()
        display_trajectory.trajectory.append(plan)

        if DEBUG == True:
            print("============ Press `Enter` to execute movement to desired square ...")
            raw_input()

        self.move_group.execute(plan, wait=True)
        # Calling `stop()` ensures that there is no residual movement
        self.move_group.stop()
        # It is always good to clear your targets after planning with poses.
        # Note: there is no equivalent function for clear_joint_value_targets()
        self.move_group.clear_pose_targets()

        return plan


    def pickAndPlace(self, action, height=0.23):
        # Save joint pose before pick&place
        self.move_group.remember_joint_values("last")
        pose_goal = self.move_group.get_current_pose().pose

        if DEBUG:
            print("Position before movement: ")
            print(pose_goal)

        # First move down (z)
        if action == "pick":
            pose_goal.position.z = height
            c = '1'
        elif action == "place":
            pose_goal.position.z = height + 0.001
            c = '0'
        else:
            return 0

        if DEBUG:
            print("Planned position: ")
            print(pose_goal)

        self.move_group.set_pose_target(pose_goal)

        plan = self.move_group.plan()

        display_trajectory = moveit_msgs.msg.DisplayTrajectory()
        display_trajectory.trajectory_start = self.robot.get_current_state()
        display_trajectory.trajectory.append(plan)

        if DEBUG == True:
            print("============ Press `Enter` to execute the movement down ...")
            raw_input()

        self.move_group.execute(plan, wait=True)
        # Calling `stop()` ensures that there is no residual movement
        self.move_group.stop()
        # It is always good to clear your targets after planning with poses.
        # Note: there is no equivalent function for clear_joint_value_targets()
        self.move_group.clear_pose_targets()

        if DEBUG:
            print("Actual: ")
            print(self.move_group.get_current_pose().pose)

        # Waiting 1 sec, control gripper and wait again
        rospy.sleep(0.1)
        self.set_gripper(c)
        rospy.sleep(0.1)

        joint_goal = self.move_group.get_named_target_values("last")

        if DEBUG == True:
            print("============ Press `Enter` to execute the movement up...")
            raw_input()

        self.move_group.go(joint_goal, wait=True)
        self.move_group.stop()

        self.move_group.forget_joint_values("last")

        return plan

    
    def dropPiece(self):
        if not self.goToStaticPosition("drop"):
            return False

        self.set_gripper(io.OFF)
        return True

    
    def pickKing(self):
        if self.king_count > 5:
            return False

        pos_name = "king" + str(self.king_count)
        if not self.goToStaticPosition(pos_name):
            return False

        # Pick king at z = 0.215
        self.pickAndPlace("pick", 0.215)

        # Increase king count
        self.king_count += 1

        return True


    def initStaticPositions(self):
        # Setting default position as move goal
        self.static_positions = yaml_handle.load_static_positions()

        return


    def goToStaticPosition(self, pose_name):
        if pose_name not in self.static_positions:
            return False

        pose_goal = self.static_positions.get(pose_name)

        self.move_group.set_pose_target(pose_goal)

        plan = self.move_group.plan()

        display_trajectory = moveit_msgs.msg.DisplayTrajectory()
        display_trajectory.trajectory_start = self.robot.get_current_state()
        display_trajectory.trajectory.append(plan)

        if DEBUG == True:
            print("============ Press `Enter` to execute movement to desired square ...")
            raw_input()

        self.move_group.execute(plan, wait=True)
        # Calling `stop()` ensures that there is no residual movement
        self.move_group.stop()
        # It is always good to clear your targets after planning with poses.
        # Note: there is no equivalent function for clear_joint_value_targets()
        self.move_group.clear_pose_targets()
        
        return True


    def initJointStaticPositions(self):
        # Setting default position as move goal
        positions = yaml_handle.load_static_positions()

        for item in positions:
            self.move_group.remember_joint_values(item, positions.get(item).pop())

        return


    def goToJointStaticPosition(self, position):
        if position not in self.move_group.get_remembered_joint_values():
            return False

        joint_goal = self.move_group.get_named_target_values(position)

        if DEBUG == True:
            print("============ Press `Enter` to execute movement to the static position...")
            raw_input()

        self.move_group.go(joint_goal, wait=True)
        self.move_group.stop()
        
        return True


    # Function for moving robot to desired locations
    def goToPoseGoal(self):
        pose_goal = geometry_msgs.msg.Pose()
        pose_goal.position.x = input("Insert x position: ")
        pose_goal.position.y = input("Insert y position: ")
        pose_goal.position.z = input("Insert z position: ")

        # Setting orientation
        pose_goal.orientation.x = 1

        if DEBUG: print(self.move_group.get_current_pose())

        self.move_group.set_pose_target(pose_goal)

        plan = self.move_group.plan()

        display_trajectory = moveit_msgs.msg.DisplayTrajectory()
        display_trajectory.trajectory_start = self.robot.get_current_state()
        display_trajectory.trajectory.append(plan)

        if DEBUG: print("============ Press `Enter` to execute movement to desired position ...")
        raw_input()

        self.move_group.execute(plan, wait=True)
        # Calling `stop()` ensures that there is no residual movement
        self.move_group.stop()
        # It is always good to clear your targets after planning with poses.
        # Note: there is no equivalent function for clear_joint_value_targets()
        self.move_group.clear_pose_targets()

        return plan

    # Function to operate gripper
    def set_gripper(self, state):
        if (state != '0') & (state != '1'):
            print "IO ERROR: Invalid gripper state"
            return False

        msg = io.get_msg(io.RDO, 8, state)

        self.io_control.publish(msg)
    