# General utilities used in multiple envs
import numpy as np
import time

from manipulator_learning.sim.utils import transformations as tf


class KeyMouseSteer(object):
    SPACEBAR = 32

    def __init__(self, pybullet_client, key_rotation = None):
        self._pb_client = pybullet_client
        self.prev_mouse_pos = np.zeros([2])
        self.mouse_pos = np.zeros(2)
        self.mouse_pos_change_rate = 3 # pixels per degree, higher --> less sensitive
        self.mouse_pos_change_rate = self.mouse_pos_change_rate * 180 / np.pi
        self.left_mb = False
        self.right_mb = False
        self.m_toggle = False
        self.d_pressed = False  # must be toggled off by user
        self.enter_toggle = False
        self.enter_hold = False
        self.enter_pressed_time = None
        self.space_toggle = False  # toggles on/off when pressed
        self.s_toggle = False
        self.policy_play = False
        self.forward_back = 0
        self.left_right = 0
        self.up_down = 0

        if key_rotation == None:
            self.key_rotation = np.eye(4)
        else:
            self.key_rotation = tf.euler_matrix(*key_rotation, axes='sxyz')

    def update_mouse(self):
        mouse = self._pb_client.getMouseEvents()
        self.prev_mouse_pos = np.copy(self.mouse_pos)

        # update mouse states for gripper joint state
        # the button states don't appear to match the flags from the documentation, it appears to be:
        #     3 - Clicked but not released since last call to getMouseEvents()
        #     4 - Released since last call to getMouseEvents()
        #     6 - Clicked and released since last call to getMouseEvents()
        for e in mouse:
            if e[0] == 1:  # mouse move event
                self.mouse_pos[0] = e[1]
                self.mouse_pos[1] = e[2]
            if e[0] == 2:  # mouse button event
                if e[3] == 0:  # left mouse button
                    if e[4] == 3:
                        self.left_mb = True
                    elif e[4] == 4:
                        self.left_mb = False
                if e[3] == 2:  # right mouse button
                    if e[4] == 3:
                        self.right_mb = True
                    elif e[4] == 4:
                        self.right_mb = False

    def update_keys(self):
        self.left_right = 0; self.forward_back = 0; self.up_down = 0
        keys = self._pb_client.getKeyboardEvents()
        for k in keys:
            if k == ord('l'):
                self.left_right = -1
            elif k == ord('j'):
                self.left_right = 1
            elif k == ord('i'):
                self.forward_back = 1
            elif k == ord('k'):
                self.forward_back = -1
            elif k == ord('o'):
                self.up_down = 1
            elif k == ord('u'):
                self.up_down = -1
            elif k == ord('m') and keys[k] == self._pb_client.KEY_WAS_RELEASED:
                if not self.m_toggle:
                    print("Mouse control on.")
                    self.m_toggle = True
                else:
                    print("Stopping mouse control.")
                    self.m_toggle = False
            elif k == KeyMouseSteer.SPACEBAR and keys[k] == self._pb_client.KEY_WAS_RELEASED:
                if not self.space_toggle:
                    self.space_toggle = True
                else:
                    self.space_toggle = False
            elif k == self._pb_client.B3G_RETURN and keys[k] == self._pb_client.KEY_WAS_RELEASED:
                if not self.enter_toggle:
                    self.enter_toggle = True
                else:
                    self.enter_toggle = False
            elif k == ord('d') and keys[k] == self._pb_client.KEY_WAS_RELEASED:
                self.d_pressed = True
            elif k == ord('s') and keys[k] == self._pb_client.KEY_WAS_RELEASED:
                if not self.s_toggle:
                    self.s_toggle = True
                else:
                    self.s_toggle = False
            elif k == ord('p') and keys[k] == self._pb_client.KEY_WAS_RELEASED:
                if not self.policy_play:
                    self.policy_play = True
                else:
                    self.policy_play = False

            if k == self._pb_client.B3G_RETURN and keys[k] == 3:
                self.enter_pressed_time = time.time()
            if self.enter_pressed_time is not None and k == self._pb_client.B3G_RETURN and \
                    (time.time() - self.enter_pressed_time > 2):
                self.enter_hold = True
            else:
                self.enter_hold = False

        # check if both keys are pressed
        if all(k in keys for k in [ord('j'), ord('l')]):  # if left and right pressed
            if keys[ord('j')] == self._pb_client.KEY_IS_DOWN and keys[ord('l')] == self._pb_client.KEY_IS_DOWN:
                self.left_right = 0
        if all(k in keys for k in [ord('i'), ord('k')]):  # if forward and backward pressed
            if keys[ord('i')] == self._pb_client.KEY_IS_DOWN and keys[ord('k')] == self._pb_client.KEY_IS_DOWN:
                self.forward_back = 0
        if all(k in keys for k in [ord('u'), ord('o')]):  # if up and down pressed
            if keys[ord('u')] == self._pb_client.KEY_IS_DOWN and keys[ord('o')] == self._pb_client.KEY_IS_DOWN:
                self.up_down = 0

    def move_gripper(self,  gripper):
        """ Return changes in a gripper state based on key/mouse commands. """

        # didn't need to redefine these, but did it mostly as a reminder that gripper is NOT
        # PASSED AS A REFERENCE
        t_control_style = gripper.t_control_style
        r_control_style = gripper.r_control_style
        g_control_style = gripper.g_control_style
        trans_acc = gripper.trans_acc
        trans_vel = gripper.trans_vel
        rot_acc = gripper.rot_acc
        rot_vel = gripper.rot_vel
        g_command = gripper.g_command
        trans_modify = None
        if r_control_style == 'p':
            quat_modify = [1, 0, 0, 0]
        else:
            quat_modify = None
        g_modify = None

        # change gripper state - b control only for now
        if g_control_style == 'b':
            g_command = self.left_mb
        # else:  # todo: implement gripper position control somehow
        #     g_command = position of some kind

        # stop if neither key is pressed
        # if self.forward_back == 0:
        #     trans_vel[0] = 0
        # if self.left_right == 0:
        #     trans_vel[1] = 0
        # if self.up_down == 0:
        #     trans_vel[2] = 0

        t_perm_change = np.dot(self.key_rotation,
                               [self.forward_back, self.left_right, self.up_down, 1])
        self.forward_back = t_perm_change[0]
        self.left_right = t_perm_change[1]
        self.up_down = t_perm_change[2]

        # GRADUALLY stop based on acceleration limits if neither key is pressed
        eps = gripper.max_trans_acc
        if self.forward_back == 0:
            if trans_vel[0] > eps:
                self.forward_back = -1
            elif trans_vel[0] < -eps:
                self.forward_back = 1
            else:
                trans_vel[0] = 0
        if self.left_right == 0:
            if trans_vel[1] > eps:
                self.left_right = -1
            elif trans_vel[1] < -eps:
                self.left_right = 1
            else:
                trans_vel[1] = 0
        if self.up_down == 0:
            if trans_vel[2] > eps:
                self.up_down = -1
            elif trans_vel[2] < -eps:
                self.up_down = 1
            else:
                trans_vel[2] = 0

        # rotate if mouse is toggled
        if self.m_toggle:
            rot_modify = [0, 0, 0]
            angle_change = (self.mouse_pos - self.prev_mouse_pos) / self.mouse_pos_change_rate
            if r_control_style == 'p':
                if not self.right_mb:  # right mb allows modifying x rotation
                    rot_modify = [0, angle_change[1], -angle_change[0]]
                else:
                    rot_modify = [angle_change[0], angle_change[1], 0]
                quat_modify = tf.quaternion_from_euler(rot_modify[0], rot_modify[1], rot_modify[2], axes='sxyz')
            elif r_control_style == 'a':
                # todo a rot won't work as is, so it should be fixed when necessary
                rot_acc = np.array(rot_modify)  # more likely to work better with a joystick/keyboard
            elif r_control_style == 'v':
                vel_modifier = 100
                if not self.right_mb:
                    rot_modify = np.array([0, angle_change[1], -angle_change[0]]) * vel_modifier
                else:
                    rot_modify = np.array([angle_change[0], angle_change[1], 0]) * vel_modifier
                rot_vel = np.array(rot_modify)

        if t_control_style == 'a':
            trans_acc = np.array([self.forward_back, self.left_right, self.up_down])
        elif t_control_style == 'v':
            mv = gripper.max_trans_vel
            trans_vel = np.array([self.forward_back*mv, self.left_right*mv, self.up_down*mv])

        return trans_acc, trans_vel, rot_acc, rot_vel, g_command, trans_modify, quat_modify, g_modify

    def move_robot(self, max_trans_vel, max_rot_vel, control_method):
        """
        Move a robot's end effector using the keyboard.

        :param max_trans_vel: Max translational velocity in any particular dimension, not total, in m/s.
        :param max_rot_vel:  Max rotational velocity in any particular dimension, rad/s.
        :param control_method: Control method being used for a Manipulator object, 'p' or 'v' for now.
        :return:
        """
        trans_vel = [0, 0, 0]
        rot_vel = [0, 0, 0]
        g_command = self.left_mb

        t_perm_change = np.dot(self.key_rotation,
                               [self.forward_back, self.left_right, self.up_down, 1])
        self.forward_back = t_perm_change[0]
        self.left_right = t_perm_change[1]
        self.up_down = t_perm_change[2]



        if self.m_toggle:
            rot_modify = [0,0,0]
            angle_change = (self.mouse_pos - self.prev_mouse_pos) / self.mouse_pos_change_rate

            if control_method == 'p':
                if self.right_mb:  # right mb allows modifying x rotation
                    rot_modify = [0, angle_change[1], -angle_change[0]]
                else:
                    rot_modify = [angle_change[0], angle_change[1], 0]
                quat_modify = tf.quaternion_from_euler(rot_modify[0], rot_modify[1], rot_modify[2], axes='sxyz')
            elif control_method == 'v':
                vel_modifier = 100
                if self.right_mb:
                    rot_modify = np.array([0, angle_change[1], angle_change[0]]) * vel_modifier
                else:
                    rot_modify = np.array([-angle_change[1], angle_change[0], 0]) * vel_modifier
                rot_vel = np.array(np.clip(rot_modify, -max_rot_vel, max_rot_vel))

        if control_method == 'v':
            trans_vel = np.array([-self.left_right * max_trans_vel, -self.up_down * max_trans_vel,
                                  self.forward_back * max_trans_vel])

        if control_method == 'p':
            raise NotImplementedError("Not yet implemented with keyboard control...")
        elif control_method == 'v':
            return trans_vel, rot_vel, g_command