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

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)
    for k in keys:
        if k == ord('l'):
            result = np.array([0.1, 0, 0, 0, 0, 0])
        elif k == ord('j'):
            result = np.array([-0.1, 0, 0, 0, 0, 0])
        elif k == ord('i'):
            result = np.array([0, 0.1, 0, 0, 0, 0])
        elif k == ord('k'):
            result = np.array([0, -0.1, 0, 0, 0, 0])
        elif k == ord('o'):
            result = np.array([0, 0, 0.1, 0, 0, 0])
        elif k == ord('u'):
            result = np.array([0, 0, -0.1, 0, 0, 0])
        elif k == ord('e'):
            result = np.array([0, 0, 0, 2, 0, 0])
        elif k == ord('q'):
            result = np.array([0, 0, 0, -2, 0, 0])
    return result

cid = p.connect(p.GUI)  # 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_SHADOWS, 1)
p.configureDebugVisualizer(p.COV_ENABLE_RENDERING, 1)
p.setGravity(0, 0, -10)
p.setRealTimeSimulation(1)

# Spawn the UR10 manipulator
m_base_pos = [1, 0, 0.9]
m_base_rot = [0, 0, 0, 1]#p.getQuaternionFromEuler([0, 0, 0])
m_ee_ind = 8
m = Manipulator(currentdir + "/../models/urdf/ur10.urdf", cid, m_ee_ind, 'p', [], m_base_pos, m_base_rot)
print(m._num_jnt_arm)
print(m._num_jnt_gripper)
create_environment()

jointPositions = [0, -3.14/2, 3.14 / 2, -3.14 / 2, -3.14 / 2, 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_ee_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

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

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

    if(cmd.any()):
        cmd[3:6] = SO3.from_quaternion(cur_rot,'xyzw').dot(cmd[3:6])
        m.set_frame_velocity_goal(m_ee_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_ee_ind , np.concatenate((w_t, w_r)), [0, 1, 2, 3, 4, 5])

    #m.close_gripper()
    m.update()
    time.sleep(0.01)

p.disconnect()