#!/usr/bin/env python
DEBUG = True
VERBOSE = True

## Python libraries
from copy import deepcopy
import sys

## 3rd party libraries
from pynput import keyboard

# Scipy (1.2.3)
from scipy.spatial.transform import Rotation

# Numpy (1.16.6)
import numpy as np

# ROS Libraries
import rospy
from geometry_msgs.msg import PoseStamped
from moveit_msgs.srv import GetPositionFK,GetPositionFKRequest
from moveit_msgs.srv import GetPositionIK,GetPositionIKRequest
from moveit_msgs.msg import MoveItErrorCodes

## Our libraries
from Controller import Controller,getR_ET


class KeyboardController(Controller):
    """
    Controller to control robot arms by keyboard.

    Most functions are implemented for convenience, not speed.
    """

    def __init__(self,
        homeConf_deg=np.array([-90,-90,45,-45,-90,-45]),
        waitForExtruder=True,
        timeScale=1.0,
        disableExtruder=False,
        robot=0):
        """
        """
        if DEBUG:
            print("Creating Controller")

        super(KeyboardController,self).__init__(
            timeScale=timeScale,
            blocking=waitForExtruder,
            disable_extruder=disableExtruder,
            robot=robot)

        # 1x6
        homeConf_rad = np.deg2rad(homeConf_deg)
        self.homeConf_rad = homeConf_rad.tolist()


    def fk(self,conf):
        """
        Compute forward kinematics using a ROS service.

        Joint angles in configurations should be in the order of getJointNames()
        Function for convenience, not for speed.

        conf: 6x numpy array of joint angles in radians
        Returns (pos_RE,quat_RE)
        pos_RE: 3x numpy array in metres
        quat_RE: 4x numpy array
        """
        if DEBUG:
            print('Controller::fk()')

        rospy.wait_for_service('compute_fk')
        try:
            moveitFK = rospy.ServiceProxy('compute_fk',GetPositionFK)
        except rospy.ServiceException as e:
            rospy.logerr("Service call failed: %s"%e)

        # Form service request
        req = GetPositionFKRequest()
        req.header.frame_id = 'base_link'
        req.fk_link_names = ['ee_link']
        req.robot_state.joint_state.name = self.group.get_active_joints()
        req.robot_state.joint_state.position = conf.tolist()

        resp = moveitFK.call(req)
        if resp.error_code.val != MoveItErrorCodes.SUCCESS:
            rospy.logwarn('Controller::fk(): Service returned {}'.format(resp.error_code.val))
            return None

        pos_RE = np.array([resp.pose_stamped.pose.position.x,resp.pose_stamped.pose.position.y,resp.pose_stamped.pose.position.z])
        quat_RE = np.array([resp.pose_stamped.pose.orientation.x,resp.pose_stamped.pose.orientation.y,resp.pose_stamped.pose.orientation.z,resp.pose_stamped.pose.orientation.w])

        return (pos_RE,quat_RE)


    def ik(self,pos_RE,quat_RE,conf_hint = None):
        """
        Compute inverse kinematics using a ROS service.
        Gives configuration that is "close" to conf_hint if specified.

        Joint angles in configurations should be in the order of getJointNames()
        Function for convenience, not for speed.
        Note: ROS service does inverse kinematics with randomness,
              so the same pose will give different joint solutions.

        pos_RE: 3x numpy array in metres
        quat_RE: 4x numpy array
        conf_hint: 6x numpy array of joint angles in radians
        Returns 6x numpy array of joint angles in radians
        """
        if DEBUG:
            print('Controller::ik()')

        if conf_hint == None:
            conf_hint = np.array(self.group.get_current_joint_values())

        rospy.wait_for_service('compute_ik')
        try:
            moveitIK = rospy.ServiceProxy('compute_ik',GetPositionIK)
        except rospy.ServiceException as e:
            rospy.logerr("Service call failed: %s"%e)


        name_conf = self.group.get_active_joints()

        req = GetPositionIKRequest()
        req.ik_request.ik_link_name = 'ee_link'
        req.ik_request.group_name = 'manipulator'
        req.ik_request.robot_state.joint_state.name = name_conf
        req.ik_request.robot_state.joint_state.position = conf_hint.tolist()
        req.ik_request.pose_stamped.pose.position.x = pos_RE[0].item()
        req.ik_request.pose_stamped.pose.position.y = pos_RE[1].item()
        req.ik_request.pose_stamped.pose.position.z = pos_RE[2].item()
        req.ik_request.pose_stamped.pose.orientation.x = quat_RE[0].item()
        req.ik_request.pose_stamped.pose.orientation.y = quat_RE[1].item()
        req.ik_request.pose_stamped.pose.orientation.z = quat_RE[2].item()
        req.ik_request.pose_stamped.pose.orientation.w = quat_RE[3].item()

        resp = moveitIK.call(req)
        if resp.error_code.val == MoveItErrorCodes.NO_IK_SOLUTION:
            return None
        elif resp.error_code.val != MoveItErrorCodes.SUCCESS:
            import pdb; pdb.set_trace()
            rospy.logwarn('Controller::ik(): Service returned {}'.format(resp.error_code.val))
            return None

        conf = np.array(resp.solution.joint_state.position)

        return conf


    def move_home(self):
        """
        Move the arm to the home configuration.
        """
        self.move_to(self.homeConf_rad,
            near_deg=None)


    def _translate_endpoint_base(self,dx_m,dy_m,dz_m,resolution_m=0.01):
        """
        Translate the endpoint of the arm from where it currently is,
        relative to the robot base.

        dx_m: 
        dy_m: 
        dz_m: 
        resolution_m: Largest end effector distance between two configurations
        """
        ## Get current pose
        pose_curr = self.group.get_current_pose().pose
        pose_next = deepcopy(pose_curr)

        ## Shift pose
        pose_next.position.x += dx_m
        pose_next.position.y += dy_m
        pose_next.position.z += dz_m

        ## Compute plan to move the end effector in a "straight" line
        (plan,frac) = self.group.compute_cartesian_path([pose_curr,pose_next],resolution_m,0.0)
        # NOTE: This plan gives duplicate points in the joint trajectory, so...

        # Create new plan with duplicates removed
        cleanTraj = deepcopy(plan.joint_trajectory)
        cleanTraj.points = []

        prevTime = None
        for point in plan.joint_trajectory.points:
            if prevTime is None or point.time_from_start.secs != prevTime.secs or point.time_from_start.nsecs != prevTime.nsecs:
                cleanTraj.points.append(point)
                prevTime = point.time_from_start

        ## Execute plan
        self.exec_traj(cleanTraj)


    def _translate_endpoint_ee(self,dx_m,dy_m,dz_m,resolution_m=0.01):
        """
        Translate the endpoint of the arm from where it currently is,
        relative to the end effector.

        dx_m: 
        dy_m: 
        dz_m: 
        resolution_m: Largest end effector distance between two configurations
        """
        ## Get current pose
        pose_curr = self.group.get_current_pose().pose

        ## Shift pose
        # Rotate translation to end effector frame
        trans_ee = np.array([dx_m,dy_m,dz_m])
        quat_curr = np.array([pose_curr.orientation.x,pose_curr.orientation.y,pose_curr.orientation.z,pose_curr.orientation.w])
        R_curr = Rotation.from_quat(quat_curr)
        trans_base = R_curr.apply(trans_ee)

        # Translate as defined at base
        self._translate_endpoint_base(trans_base[0].item(),trans_base[1].item(),trans_base[2].item(),resolution_m)


    def _translate_endpoint_tool(self,dx_m,dy_m,dz_m,resolution_m=0.01):
        """
        Translate the endpoint of the arm from where it currently is,
        relative to the tool tip.

        dx_m: 
        dy_m: 
        dz_m: 
        resolution_m: Largest end effector distance between two configurations
        """
        ## Get orientation of tool tip relative to end effector
        R_ET = getR_ET()

        ## Shift pose
        # Rotate translation to end effector frame
        trans_tool = np.array([dx_m,dy_m,dz_m])
        trans_ee = R_ET.apply(trans_tool)

        # Translate as defined at end effector
        self._translate_endpoint_ee(trans_ee[0].item(),trans_ee[1].item(),trans_ee[2].item(),resolution_m)


    def _rotate_endpoint_base(self,dx_deg,dy_deg,dz_deg):
        """
        Rotate the endpoint of the arm from where it currently is,
        relative to the robot base.

        Rotation is specified as a rotation vector
        See: https://en.wikipedia.org/wiki/Axis%E2%80%93angle_representation#Rotation_vector

        dx_deg:
        dy_deg:
        dz_deg:
        Returns True if arm is moved, False otherwise
        """
        ## Get current pose
        pose_curr = self.group.get_current_pose().pose

        ## Get current orientation
        # 4x
        quat_curr = np.array([pose_curr.orientation.x,pose_curr.orientation.y,pose_curr.orientation.z,pose_curr.orientation.w])

        R_curr = Rotation.from_quat(quat_curr)

        ## Get step rotation
        # 3x
        rotVec_base_deg = np.array([dx_deg,dy_deg,dz_deg])
        rotVec_base_rad = np.deg2rad(rotVec_base_deg)

        ## Calculate amount of rotation in end effector frame
        R_step_base = Rotation.from_rotvec(rotVec_base_rad)

        R_step_ee = R_curr.inv() * R_step_base * R_curr;

        rotVec_ee_rad = R_step_ee.as_rotvec()
        rotVec_ee_deg = np.rad2deg(rotVec_ee_rad)

        ## Rotate as defined at end effector
        return self._rotate_endpoint_ee(rotVec_ee_deg[0].item(),rotVec_ee_deg[1].item(),rotVec_ee_deg[2].item())


    def _rotate_endpoint_ee(self,dx_deg,dy_deg,dz_deg):
        """
        Rotate the endpoint of the arm from where it currently is, relative to the end effector

        Arguments as rotation vector
        See: https://en.wikipedia.org/wiki/Axis%E2%80%93angle_representation#Rotation_vector

        dx:
        dy:
        dz:
        Returns True if arm is moved, False otherwise
        """
        # Number of intermediate steps to get to the final configuration

        ## Get current pose
        pose_curr = self.group.get_current_pose().pose

        ## Get current orientation
        # 4x
        quat_curr = np.array([pose_curr.orientation.x,pose_curr.orientation.y,pose_curr.orientation.z,pose_curr.orientation.w])

        R_curr = Rotation.from_quat(quat_curr)

        ## Get step rotation
        # 3x
        rotVec_deg = np.array([dx_deg,dy_deg,dz_deg])
        rotVec_rad = np.deg2rad(rotVec_deg)

        ## Calculate intermediate shift orientation
        R_step = Rotation.from_rotvec(rotVec_rad)

        R_next = R_curr * R_step;

        quat_next = R_next.as_quat()
        pos_next = np.array([pose_curr.position.x,pose_curr.position.y,pose_curr.position.z])

        ## Convert to joint angles
        conf_next = self.ik(pos_next,quat_next)
        if conf_next is None:
            return False

        ## Visualise plan to confirm
        # TODO

        ## Execute plan
        success = self.move_to(conf_next,
            near_deg=2*np.linalg.norm(rotVec_deg))

        return success


    def _move_callback(self,key,configurations,params):
        """
        Callback function to operate the robot arm.

        Relative to robot base:
        - the robot can be translated in the:
            - XY plane with WASD
            - Z direction with left shift and left ctrl
        - the robot can be rotated in the:
            - XY directions with IJKL
            - Z direction with U and O

        Relative to end effector:
        - the robot can be translated in the:
            - XY plane with arrow keys
            - Z direction with right shift and right ctrl
        - the robot can be rotated in the:
            - XY directions with the 8456 keys (numpad)
            - Z direction with 7 and 9

        The -/_ and =/+ keys will adjust the amount to translate and rotate.

        The page up/down keys will retract and load filament from the print tip.

        Enter key will save the joint angles and pose.
        Esc key will stop the gathering of configurations.
        Home key will move the robot to the home configuration.
        """
        if DEBUG:
            print('Controller::_move_callback()')

        step_m = params['step_m']
        step_deg = min(params['step_deg'],45)

        ## Base frame controls
        # Translation
        if key in (keyboard.KeyCode.from_char('w'),keyboard.KeyCode.from_char('W')):
            if VERBOSE:
                print('Step {} mm in y-direction wrt robot base'.format(-step_m*1000))
            self._translate_endpoint_base(0,-step_m,0,step_m/100)
        elif key in (keyboard.KeyCode.from_char('s'),keyboard.KeyCode.from_char('S')):
            if VERBOSE:
                print('Step {} mm in y-direction wrt robot base'.format(step_m*1000))
            self._translate_endpoint_base(0,step_m,0,step_m/100)
        elif key in (keyboard.KeyCode.from_char('a'),keyboard.KeyCode.from_char('A')):
            if VERBOSE:
                print('Step {} mm in x-direction wrt robot base'.format(step_m*1000))
            self._translate_endpoint_base(step_m,0,0,step_m/100)
        elif key in (keyboard.KeyCode.from_char('d'),keyboard.KeyCode.from_char('D')):
            if VERBOSE:
                print('Step {} mm in x-direction wrt robot base'.format(-step_m*1000))
            self._translate_endpoint_base(-step_m,0,0,step_m/100)

        elif key == keyboard.Key.shift_l:
            if VERBOSE:
                print('Step {} mm in z-direction wrt robot base'.format(step_m*1000))
            self._translate_endpoint_base(0,0,step_m,step_m/100)
        elif key == keyboard.Key.ctrl_l:
            if VERBOSE:
                print('Step {} mm in z-direction wrt robot base'.format(-step_m*1000))
            self._translate_endpoint_base(0,0,-step_m,step_m/100)

        # Rotation
        elif key in (keyboard.KeyCode.from_char('i'),keyboard.KeyCode.from_char('I')):
            if VERBOSE:
                print('Rotate {} deg in the x-direction wrt robot base'.format(step_deg))
            self._rotate_endpoint_base(step_deg,0,0)
        elif key in (keyboard.KeyCode.from_char('k'),keyboard.KeyCode.from_char('K')):
            if VERBOSE:
                print('Rotate {} deg in the x-direction wrt robot base'.format(-step_deg))
            self._rotate_endpoint_base(-step_deg,0,0)
        elif key in (keyboard.KeyCode.from_char('j'),keyboard.KeyCode.from_char('J')):
            if VERBOSE:
                print('Rotate {} deg in the y-direction wrt robot base'.format(step_deg))
            self._rotate_endpoint_base(0,step_deg,0)
        elif key in (keyboard.KeyCode.from_char('l'),keyboard.KeyCode.from_char('L')):
            if VERBOSE:
                print('Rotate {} deg in the y-direction wrt robot base'.format(-step_deg))
            self._rotate_endpoint_base(0,-step_deg,0)

        elif key in (keyboard.KeyCode.from_char('u'),keyboard.KeyCode.from_char('U')):
            if VERBOSE:
                print('Rotate {} deg in the z-direction wrt robot base'.format(step_deg))
            self._rotate_endpoint_base(0,0,step_deg)
        elif key in (keyboard.KeyCode.from_char('o'),keyboard.KeyCode.from_char('O')):
            if VERBOSE:
                print('Rotate {} deg in the z-direction wrt robot base'.format(-step_deg))
            self._rotate_endpoint_base(0,0,-step_deg)

        ## End effector frame controls
        # Translation
        elif key == keyboard.Key.up:
            if VERBOSE:
                print('Step {} mm in z-direction wrt end effector'.format(step_m*1000))
            self._translate_endpoint_ee(0,0,step_m,step_m/100)
        elif key == keyboard.Key.down:
            if VERBOSE:
                print('Step {} mm in z-direction wrt end effector'.format(-step_m*1000))
            self._translate_endpoint_ee(0,0,-step_m,step_m/100)
        elif key == keyboard.Key.left:
            if VERBOSE:
                print('Step {} mm in y-direction wrt end effector'.format(-step_m*1000))
            self._translate_endpoint_ee(0,step_m,0,step_m/100)
        elif key == keyboard.Key.right:
            if VERBOSE:
                print('Step {} mm in y-direction wrt end effector'.format(step_m*1000))
            self._translate_endpoint_ee(0,-step_m,0,step_m/100)

        elif key == keyboard.Key.shift_r:
            if VERBOSE:
                print('Step {} mm in x-direction wrt end effector'.format(-step_m*1000))
            self._translate_endpoint_ee(-step_m,0,0,step_m/100)
        elif key == keyboard.Key.ctrl_r:
            if VERBOSE:
                print('Step {} mm in x-direction wrt end effector'.format(step_m*1000))
            self._translate_endpoint_ee(step_m,0,0,step_m/100)

        # Rotation
        elif key == keyboard.KeyCode.from_char('8'):
            if VERBOSE:
                print('Rotate {} deg in the y-direction wrt end effector'.format(step_deg))
            self._rotate_endpoint_ee(0,step_deg,0)
        elif key == keyboard.KeyCode.from_char('5') or (hasattr(key,'vk') and key.vk is not None and int(key.vk) == 65437):
            if VERBOSE:
                print('Rotate {} deg in the y-direction wrt end effector'.format(-step_deg))
            self._rotate_endpoint_ee(0,-step_deg,0)
        elif key == keyboard.KeyCode.from_char('4'):
            if VERBOSE:
                print('Rotate {} deg in the z-direction wrt end effector'.format(-step_deg))
            self._rotate_endpoint_ee(0,0,-step_deg)
        elif key == keyboard.KeyCode.from_char('6'):
            if VERBOSE:
                print('Rotate {} deg in the z-direction wrt end effector'.format(step_deg))
            self._rotate_endpoint_ee(0,0,step_deg)

        elif key == keyboard.KeyCode.from_char('7'):
            if VERBOSE:
                print('Rotate {} deg in the x-direction wrt end effector'.format(-step_deg))
            self._rotate_endpoint_ee(-step_deg,0,0)
        elif key == keyboard.KeyCode.from_char('9'):
            if VERBOSE:
                print('Rotate {} deg in the x-direction wrt end effector'.format(step_deg))
            self._rotate_endpoint_ee(step_deg,0,0)

        elif key in (keyboard.KeyCode.from_char('-'),keyboard.KeyCode.from_char('_')): # Decrease step resolution
            if params['step_m']/2 != 0:
                params['step_m'] = params['step_m']/2
            if params['step_deg']/2 != 0:
                params['step_deg'] = params['step_deg']/2
            if VERBOSE:
                print('Decreased step size to {} mm/{} deg'.format(params['step_m']*1000,params['step_deg']))
        elif key in (keyboard.KeyCode.from_char('+'),keyboard.KeyCode.from_char('=')): # Increase step resolution
            if params['step_m']*2 != float('inf'):
                params['step_m'] = params['step_m']*2
            if params['step_deg']*2 != float('inf'):
                params['step_deg'] = params['step_deg']*2
            if VERBOSE:
                print('Increased step size to {} mm/{} deg'.format(params['step_m']*1000,params['step_deg']))

        elif key == keyboard.Key.page_up:
            if VERBOSE:
                print('Retract')
            self.postprint_retract()
        elif key == keyboard.Key.page_down:
            if VERBOSE:
                print('Extrude')
            self.preprint_load()

        elif key == keyboard.Key.home:
            if VERBOSE:
                print('Go home')
            self.move_home()
        elif key == keyboard.Key.enter:
            if VERBOSE:
                print('Saving configuration')
            params['success'] = True
            configurations.append(
                (self.group.get_current_joint_values(),
                self.group.get_current_pose()))
            if DEBUG:
                print('Returning control to program')
            return False
        elif key == keyboard.Key.esc:
            if VERBOSE:
                print('Stopping keyboard controls')
            params['success'] = False
            if DEBUG:
                print('Returning control to program')
            return False
        else:
            # Do nothing for now
            pass


    def prompt_configurations(self,prompt=None):
        """
        Prompts user to control the robot to a series of configurations.
        Relative to robot base, the robot can be translated in the:
        - XY plane with WASD
        - Z direction with shift and ctrl
        The -/_ and =/+ keys will adjust the amount to translate (default is 10mm).

        Relative to the end effector, the robot can be rotated 90 degrees in the:
        - XY directions with the arrow keys
        - Z direction with Q and E

        Enter key will save the joint angles and pose.
        Esc key will SystemExit.
        Home key will move the robot to the home configuration.

        The variable prompt can take multiple types:
        - List of strings: Prompts for the user before every configuration.
          Will stop gathering when all the prompts are addressed.
        - Positive integer: Stops gathering after a number of configurations.
        - None: Will keep gathering until Esc, which will not SystemExit.

        Returns list of (list of joint angle, ROS pose message) tuples
        """
        if DEBUG:
            print('Controller::prompt_configurations()')

        configurations = []
        params = {
            'step_m': 0.016,
            'step_deg': 11.25,
            'success': True}

        if isinstance(prompt,list):
            for p in prompt:
                print(p)
                with keyboard.Listener(
                        on_release = lambda key:self._move_callback(key,configurations,params),
                        suppress=True) as listener:
                    listener.join()
                    if DEBUG:
                        print('Keyboard listener thread joined')
                if not params['success']:
                    sys.exit(0)
        elif isinstance(prompt,int):
            for iConfig in range(prompt):
                print('Enter configuration {}'.format(iConfig))
                with keyboard.Listener(
                        on_release=lambda key:self._move_callback(key,configurations,params),
                        suppress=True) as listener:
                    listener.join()
                    if DEBUG:
                        print('Keyboard listener thread joined')
                if not params['success']:
                    sys.exit(0)
        else:
            while params['success']:
                params['success'] = False

                print('Enter next configuration')
                with keyboard.Listener(
                        on_release=lambda key:self._move_callback(key,configurations,params),
                        suppress=True) as listener:
                    listener.join()

        return configurations


if __name__=="__main__":
    #con = KeyboardController(disableExtruder=True)
    con = KeyboardController()
    #con1 = KeyboardController(disableExtruder=True, robot=1)
    #con2 = KeyboardController(disableExtruder=True, robot=2)

    print('Collecting configrations...')
    configuations = con.prompt_configurations()
    #configuations = con1.prompt_configurations()
    #configuations = con2.prompt_configurations()
    print('Your configurations:')
    print(configuations)