#! /usr/bin/env python

# *******************************************************************
# Author: Sahand Rezaei-Shoshtari
# Oct. 2019
# Copyright 2019, Sahand Rezaei-Shoshtari, All rights reserved.
# *******************************************************************

import numpy as np
import sys
import trajectory
import copy
import matplotlib.pyplot as plt

# ROS libs
import rospy
import moveit_commander
import tf.transformations

# ROS messages
from geometry_msgs.msg import Pose, PoseStamped
from sensor_msgs.msg import JointState
from moveit_msgs.msg import RobotState
from std_msgs.msg import Header


class Planner:
    """
    This class handles the motion planning of the robot.
    """
    def __init__(self, robot_name):
        """
        :param robot_name: 'j2n6s300' or 'j2s7s300'
        :type robot_name: str
        """
        self.n_joints = int(robot_name[3])

        # trajectory in joint space
        self.traj = trajectory.Trajectory(self.n_joints)

        # MoveIt! objects
        self.robot = None
        self.move_group = None
        self.scene = None

    def create_test_traj(self, pos_amp, freq, start_pos, total_t, time_step=0.01):
        """
        Creates a periodic trajectory for testing the controllers. The motion is generated by superimposing a series of
        sinusoids functions.
        :param pos_amp: amplitude of the joint positions (in Radians) - 2D array [n_func, n_joints]
        :type pos_amp: np.array
        :param freq: frequency of the periodic functions - 2D array [n_func, n_joints]
        :type freq: np.array
        :param start_pos: starting position of the trajectory (same as the current position of the robot)
        :type start_pos: np.array or list
        :param total_t: total duration of the trajectory
        :type total_t: float or int
        :param time_step: time between each waypoint
        :type time_step: float
        :return: trajectory
        :rtype: trajectory.Trajectory
        """
        t_waypoints = np.reshape(np.arange(0., total_t, time_step), (-1, 1))
        waypoints = np.zeros((len(t_waypoints), self.n_joints))
        omega = 2 * np.pi * freq

        for index, t in np.ndenumerate(t_waypoints):
            waypoints[index[0], :] = np.sum(pos_amp * np.sin(omega * t), axis=0) + np.array(start_pos)

        self.traj = trajectory.Trajectory(self.n_joints, waypoints=waypoints, t_waypoints=t_waypoints, start_pos=start_pos,
                                          step_size=time_step)

        return self.traj

    def init_moveit(self):
        """
        Initializes the inverse kinematics system. Required if inverse kinematics is going to be used.
        :return None
        """
        # set up MoveIt! robot commander and move moveit_group commander
        moveit_commander.roscpp_initialize(sys.argv)
        self.robot = moveit_commander.RobotCommander()
        self.move_group = moveit_commander.MoveGroupCommander("arm")

        # set up MoveIt! moveit_scene interface
        self.scene = moveit_commander.PlanningSceneInterface()

        # setting the position and orientation tolerance
        self.move_group.set_goal_position_tolerance(0.001)
        self.move_group.set_goal_orientation_tolerance(0.001)

        # Add a table to the moveit_scene acting as a ground
        rospy.sleep(1)
        table_position = PoseStamped()
        table_position.header.frame_id = self.robot.get_planning_frame()
        table_position.pose.position.x = 0.
        table_position.pose.position.y = 0.
        table_position.pose.position.z = -.03 / 2
        self.scene.add_box("table", table_position, (2.4, 2.4, 0.03))

        rospy.loginfo("MoveIt! init successful.")

    def plan_moveit(self, position, orientation, start_joint_position=None, euler_flag=False, time_scale=1.0):
        """
        Plans the motions using the IK implemented in MoveIt!. Note that MoveIt! is not used to execute the plan on the
        robot, because the controllers are not very useful.
        :param position: list of cartesian coordinates (XYZ)
        :type position: list or np.array
        :param orientation: list of orientation elements (default: quaternion)
        :type orientation: list or np.array
        :param start_joint_position: starting joint position for the planner
        :type: list or np.array
        :param euler_flag: a flag to indicate whether the orientation is in euler or quaternions
        :type euler_flag: bool
        :param time_scale: scales time for the interpolation of raw waypoints
        :type time_scale: float
        :return the trajectory with the raw waypoints
        :rtype: trajectory.Trajectory
        """
        if start_joint_position is not None:
            self.set_start_pose_moveit(start_joint_position)

        pose = self.create_pose(position, orientation, euler_flag=euler_flag)
        self.move_group.set_pose_target(pose)
        plan = self.move_group.plan()
        self.move_group.clear_pose_targets()

        # create the raw waypoints
        raw_waypoints = np.array([p.positions for p in plan.joint_trajectory.points])
        t_raw_waypoints = np.array([self.to_sec(p.time_from_start) for p in plan.joint_trajectory.points])

        self.traj = trajectory.Trajectory(self.n_joints, raw_waypoints=raw_waypoints, t_raw_waypoints=t_raw_waypoints,
                                          time_scale=time_scale)

        return self.traj

    def plan_cutting(self, start_position=[0., -0.4, 0.5], end_position=[0., -0.7, 0.5],
                     orientation=[0.7071, 0., 0., .7071], cutting_direction=[0., -1., 0.], cutting_plane=[1, 1, 0],
                     rod_center=[0., -0.55, 0.50], rod_radius=0.08, cut_force_k=100., cut_force_d=10., count=10):
        """
        Plans a simple cutting scenario using MoveIt! planner.
        :param start_position: starting position
        :param end_position: end position
        :param orientation: orientation during cutting
        :param cutting_direction: cutting direction
        :param cutting_plane: cutting plane
        :param rod_center: center of the rod position
        :param rod_radius: radius of the tree
        :type rod_radius: float
        :param cut_force_k: K of the cutting force
        :type cut_force_k: float         
        :param cut_force_d: D of the cutting force
        :type cut_force_d: float
        :param count: number of cuts
        :type count: int
        :return: trajectory
        :rtype: trajectory.Trajectory
        """        
        # start of the motion
        start = self.plan_moveit(position=start_position, orientation=orientation)
        start_joint_position = start.waypoints[-1]            
        
        # trajectory 1: go to the starting position
        start_motion = self.create_test_traj(pos_amp=np.array([[0.] * 6]), freq=np.array([[0.2] * 6]),
                                             start_pos=start_joint_position, total_t=5)

        # trajectory 2: forward motion
        forward_motion = self.plan_moveit(position=end_position, orientation=orientation,
                                          start_joint_position=start_joint_position, time_scale=1.)

        # trajectory 3: backward motion
        backward_motion = self.plan_moveit(position=start_position, orientation=orientation,
                                           start_joint_position=forward_motion.waypoints[-1], time_scale=1.)

        # combine all trajectories
        cutting_motion = self.append_trajectory(forward_motion, backward_motion)
        cutting_traj = self.repeat_trajectory(cutting_motion, count)

        traj = self.append_trajectory(start_motion, cutting_traj)

        # set the position and radius of the wooden rod and the cutting direction
        traj.set_interaction_force_param(rod_center, rod_radius, cut_force_k, cut_force_d, cutting_direction, cutting_plane)

        return traj

    def set_start_pose_moveit(self, joint_position):
        """
        Sets the start position of the robot for MoveIt! planner. Note that this method does not move the robot but only
        specifies the starting position for path planning.
        :param joint_position: starting position
        :type joint_position: list or np.array
        :return: None
        """
        # create the robot state msg
        joint_state = JointState()
        joint_state.header = Header()
        joint_state.header.stamp = rospy.Time.now()
        joint_state.name = ['j2n6s300_joint_1', 'j2n6s300_joint_2', 'j2n6s300_joint_3', 'j2n6s300_joint_4',
                            'j2n6s300_joint_5', 'j2n6s300_joint_6']
        joint_state.position = list(joint_position)
        moveit_robot_state = RobotState(joint_state=joint_state)

        # set the start state for the planner
        self.move_group.set_start_state(moveit_robot_state)

    def append_trajectory(self, traj_1, traj_2):
        """
        Appends traj_2 to the end of traj_1. Warning: it is up to the user to check for continuity of the trajectories.
        :param traj_1: trajectory 1
        :type traj_1: trajectory.Trajectory
        :param traj_2: trajectory 2
        :type traj_1: trajectory.Trajectory
        :return: combined trajectory
        :rtype: trajectory.Trajectory
        """
        step_size = min(traj_1.step_size, traj_2.step_size)
        waypoints = np.append(traj_1.waypoints, traj_2.waypoints, axis=0)
        t_waypoints = np.append(traj_1.t_waypoints, traj_2.t_waypoints + traj_1.total_t + step_size, axis=0)
        start_pos = traj_1.start_pos

        self.traj = trajectory.Trajectory(self.n_joints, waypoints=waypoints, t_waypoints=t_waypoints, start_pos=start_pos,
                                          step_size=step_size)
        return self.traj

    def repeat_trajectory(self, traj, repeat_count):
        """
        Repeats the trajectory by appending a duplicate to the end of it.
        :param traj: trajectory
        :type traj: trajectory.Trajectory
        :param repeat_count: number of repeats
        :type repeat_count: int
        :return: trajectory
        :rtype: trajectory.Trajectory
        """
        repeated_traj = copy.copy(traj)
        for i in range(repeat_count - 1):
            repeated_traj = self.append_trajectory(repeated_traj, traj)

        return repeated_traj

    @staticmethod
    def to_sec(duration):
        """
        Converts rospy.Duration to a second.
        :param duration: time
        :type duration: rospy.Duration
        :return: converted timm in seconds
        :rtype: float
        """
        return duration.secs + duration.nsecs * 10**-9

    @staticmethod
    def create_pose(position, orientation, euler_flag=False):
        """
        Creates a pose for the end effector.
        :param position: list of cartesian coordinates (XYZ)
        :type position: list or np.array
        :param orientation: list of orientation elements (default: quaternion)
        :type orientation: list or np.array
        :param euler_flag: a flag to indicate whether the orientation is in euler or quaternions
        :type euler_flag: bool
        :return the pose
        :rtype: Pose
        """
        # convert orientation euler angles to quaternions if the optional flag is set to True
        if euler_flag:
            orientation = tf.transformations.quaternion_from_euler(*orientation)

        # create the pose
        pose = Pose()
        pose.orientation.x = orientation[0]
        pose.orientation.y = orientation[1]
        pose.orientation.z = orientation[2]
        pose.orientation.w = orientation[3]
        pose.position.x = position[0]
        pose.position.y = position[1]
        pose.position.z = position[2]

        return pose

    @staticmethod
    def shut_down_moveit():
        """
        Shuts down MoveIt! planner.
        :return: None
        """
        rospy.loginfo("Shutting down MoveIt!")
        moveit_commander.roscpp_shutdown()


if __name__ == '__main__':
    # a simple test for checking the trajectory parametrization and differentiation
    planner = Planner()
    starting_position = [4.7, 3.5, 2., 4.7, 0., 1.57]

    # trajectory 1: go to the starting position
    test_pos = np.array([[0.0, 0.0, 0.5, 0.0, 0.0, 0.0]])
    test_freq = np.array([[0.2, 0.2, 0.2, 0.2, 0.2, 0.2]])
    test_traj = planner.create_test_traj(test_pos, test_pos, starting_position, 5, time_step=0.01)

    # combine all trajectories
    test_traj = planner.repeat_trajectory(test_traj, 5)

    plt.subplot(3, 1, 1)
    plt.plot(test_traj.t_waypoints, test_traj.waypoints[:, 2], 'g')

    plt.subplot(3, 1, 2)
    plt.plot(test_traj.t_waypoints, test_traj.joint_vel[:, 2], 'b')

    plt.subplot(3, 1, 3)
    plt.plot(test_traj.t_waypoints, test_traj.joint_acc[:, 2], 'r')

    plt.show()
