# -*- 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 numpy import linalg as la

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

    return result


def matprint(mat, fmt="g"):
    col_maxes = [max([len(("{:" + fmt + "}").format(x)) for x in col]) for col in mat.T]
    for x in mat:
        for i, y in enumerate(x):
            print(("{:" + str(col_maxes[i]) + fmt + "}").format(y), end="  ")
        print("")


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 = p.getQuaternionFromEuler([0, 0, 0])
m_ee_ind = 7
m = Manipulator(currentdir + "/../models/urdf/jaco.urdf", cid, m_ee_ind, 'p', [9,11,13], m_base_pos, m_base_rot)
create_environment()

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)
m.update()  # update joint position

# 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

m.set_control_method('v') # we want velocity control
M = m.get_manipulability_ellipsoid('force',[0,1,2])
vis = p.createVisualShape(p.GEOM_SPHERE, radius=0.01, visualFramePosition=[0, 0, 0], rgbaColor=[1, 0, 0, 1])

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()):
        goal_pos = cur_pos + cmd
    w_t, w_r = mu.pose_error(cur_pos, cur_rot, goal_pos, goal_rot)
    m.set_vel_and_mnp_goal(m_ee_ind, np.concatenate((w_t, w_r)),M,[0,1],[0,1,2])
    m.close_gripper()
    m.update()
    print(mu.SPD_error(m.get_manipulability_ellipsoid('force',[0,1,2]),M).reshape(-1))

p.disconnect()
