# -*- coding: utf-8 -*-
import os, inspect
import numpy as np
# import pybullet as p
import pybullet
import time
import pybullet_data
from sim.robots.pyb_manipulator.manipulator import Manipulator
from sim.robots.pyb_manipulator.utils import math_utils as mu
from liegroups.numpy import SO3

from sim.robots.utils.bullet_client import BulletClient


p = BulletClient(pybullet.GUI_SERVER)

def create_environment():
    plane = [p.loadURDF("plane.urdf", [0.000000, 0.000000, 0.000000], [0.000000, 0.000000, 0.000000, 1.000000])]  # load plane
    objects = [p.loadURDF("table/table.urdf", [1.000000, -0.200000, 0.000000], [0.000000, 0.000000, 0.707107, 0.707107])]

def update_keys():
    keys = p.getKeyboardEvents()
    result = np.zeros(6)
    grip = False
    for k in keys:
        if k == ord('l'):
            result = np.array([0.05, 0, 0, 0, 0, 0])
        elif k == ord('j'):
            result = np.array([-0.05, 0, 0, 0, 0, 0])
        elif k == ord('i'):
            result = np.array([0, 0.05, 0, 0, 0, 0])
        elif k == ord('k'):
            result = np.array([0, -0.05, 0, 0, 0, 0])
        elif k == ord('o'):
            result = np.array([0, 0, 0.05, 0, 0, 0])
        elif k == ord('u'):
            result = np.array([0, 0, -0.05, 0, 0, 0])
        elif k == ord('y'):
            if keys[k] == p.KEY_IS_DOWN:
                grip = True
    return result, grip

# cid = p.connect(p.GUI_SERVER)  # stuff above hangs much of the time for some reason
currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))  # get current directory

p.resetSimulation()
p.setAdditionalSearchPath(pybullet_data.getDataPath())
p.configureDebugVisualizer(p.COV_ENABLE_RENDERING, 1)
p.setGravity(0, 0, -10)
p.setRealTimeSimulation(1)

# Spawn the Jaco manipulator
m_base_pos = [1, 0, 0.9]
m_base_rot = [0, 0, 0, 1]#p.getQuaternionFromEuler([0, 0, 0])
m_ee_ind = 7
m_tool_ind = 8
m = Manipulator(p, currentdir + "/../models/urdf/jaco.urdf", m_ee_ind,
                m_tool_ind, 'p', [9,11,13], m_base_pos, m_base_rot, self_collision=True)
create_environment()

ee = m._ee_link_ind

p.addUserDebugLine([0,0,0],[0.3,0,0],[1,0,0],parentObjectUniqueId=m._arm[0],
                   parentLinkIndex=8, lineWidth=5)
p.addUserDebugLine([0,0,0],[0,0.3,0],[0,1,0],parentObjectUniqueId=m._arm[0],
                   parentLinkIndex=8, lineWidth=5)
p.addUserDebugLine([0,0,0],[0,0,0.3],[0,0,1],parentObjectUniqueId=m._arm[0],
                   parentLinkIndex=8, lineWidth=5)

p.resetDebugVisualizerCamera(cameraDistance=1.0, cameraYaw=-90, cameraPitch=-45,
                                                   cameraTargetPosition=[1, -.2, .7])

jointPositions = [0, 3.14, 3.14 / 2, -3.14 / 2, -3.14 / 2, 0, 0, 0, 0]  # set joint position goal
m.set_joint_position_goal(jointPositions)
m.update()  # update joint position
time.sleep(1)

# set up starting pose and goal
pose = m.get_link_pose(m_tool_ind)
start_pos = pose[0:3]
start_rot = pose[3:7]
goal_pos = start_pos
goal_rot = start_rot
cur_pos = start_pos
cur_rot = start_rot

cmd_given = False

m.set_control_method('v') # we want velocity control
while (1):
    cmd, grip_cmd = update_keys()
    pose = m.get_link_pose(m_tool_ind)

    cur_pos = pose[0:3]
    cur_rot = pose[3:7]

    if(cmd.any()):
        cmd_given = True
        cmd[3:6] = SO3.from_quaternion(cur_rot,'xyzw').dot(cmd[3:6])

        # add a transform here for cmd[0:3] to also be in the frame of m tool ind
        # t_mat = tf.quaternion_matrix(convert_quat_pb_to_tf(cur_rot))
        # t_mat[:3,3] = cur_pos
        # cmd[0:3] = np.dot(t_mat, np.array([*cmd[0:3], 1]))[:3]

        cmd[:3] = SO3.from_quaternion(cur_rot, 'xyzw').dot(cmd[:3])

        m.set_frame_velocity_goal(m_tool_ind, cmd, [0, 1, 2, 3, 4, 5])
        goal_pos = pose[0:3]
        goal_rot = pose[3:7]
    else:
        w_t, w_r = mu.pose_error(cur_pos, cur_rot, goal_pos, goal_rot)
        m.set_frame_velocity_goal(m_tool_ind , np.concatenate((w_t, w_r)), [0, 1, 2, 3, 4, 5])

    if grip_cmd:
        m.close_gripper()
    else:
        m.open_gripper()
    m.update()
    time.sleep(0.01)

p.disconnect()
