import gym
from gym import error, spaces, utils
from gym.utils import seeding

import os
import pybullet as p
import pybullet_data as pd
import math
import numpy as np
import random


class UREnv(gym.Env):
    metadata = {'render.modes': ['human']}

    def __init__(self):
        p.connect(p.GUI)
        # 相机位置
        # p.resetDebugVisualizerCamera()

        self.action_space = spaces.Box(np.array([-1] * 4), np.array([1] * 4))
        self.observation_space = spaces.Box(np.array([-1] * 5), np.array([1] * 5))

    def step(self, action):
        p.configureDebugVisualizer(p.COV_ENABLE_SINGLE_STEP_RENDERING)
        orientation = p.getQuaternionFromEuler([0., -math.pi, math.pi / 2.])
        dv = 0.005
        dx = action[0] * dv
        dy = action[1] * dv
        dz = action[2] * dv
        fingers = action[3]

        currentPose = p.getLinkState(self.ur5_Id, 11)
        # currentPosition = currentPOse[0]
        currentPosition = currentPose[4]
        newPosition = [currentPosition[0] + dx,
                       currentPosition[1] + dy,
                       currentPosition[2] + dz]
        jointPoses = p.calculateInverseKinematics(self.ur5_Id, 11, newPosition, orientation)
        p.stepSimulation()

        state_object, _ = p.getBasePositionAndOrientation(self.objectUid)
        state_robot = p.getLinkState(self.ur5_Id, 11)[0]
        state_fingers = (p.getJointState(self.ur5_Id, 9)[0], p.getJointState(self.ur5_Id, 10)[0])

        # 如果将物体抓住了，并且抓取到了0.45高度，会获得+1奖励
        if state_object[2] > 0.45:
            reward = 1
            done = True
        else:
            reward = 0
            done = False

        info = state_object
        observation = state_robot + state_fingers
        return observation, reward, done, info

    def reset(self):
        p.resetSimulation()
        # 关闭可视化
        p.configureDebugVisualizer(p.COV_ENABLE_RENDERING, 0)
        p.setGravity(0, 0, -10)

        p.setAdditionalSearchPath(pd.getDataPath())
        planeUid = p.loadURDF('plane.urdf', basePosition=[0, 0, -0.65])
        rest_poses = [0, -0.215, 0, -0.275, 0, 2.356, 2.356, 0.08, 0.08]
        self.ur5_Id = p.loadURDF('/Users/iub/强化学习/pybullet_ur5_robotiq-robotflow/urdf/ur5.urdf', useFixedBase=True)

        # 将机械臂关节放置设定位置
        for i in range(7):
            p.resetJointState(self.ur5_Id, rest_poses[i])

        tableUid = p.loadURDF('table.urdf', basePosition=[0.5, 0, -0.65])
        trayUid = p.loadURDF('traybox.urdf', basePosition=[0.65, 0, 0])
        state_object = [random.uniform(0.5, 0.8), random.uniform(-0.2, 0.2), 0.05]
        self.objectUid = p.loadURDF('random_urdfs/000/000.urdf', basePosition=state_object)
        state_robot = p.getLinkState(self.ur5_Id, 11)[0]
        state_fingers = (p.getJointState(self.ur5_Id, 9)[0], p.getJointState(self.ur5_Id, 10)[0])
        observation = state_robot + state_fingers
        p.configureDebugVisualizer(p.COV_ENABLE_RENDERING, 1)
        return observation

    def render(self, mode='human'):
        # 放置相机到目标位置和姿势
        view_matrix = p.computeViewMatrixFromYawPitchRoll(cameraTargetPosition=[0.7, 0, 0.5],
                                                          distance=.7,
                                                          yaw=90,
                                                          pitch=-70,
                                                          roll=0, upAxisIndex=2)
        proj_matrix = p.computeProjectionMatrixFOV(fov=60, aspect=float(960) / 720,
                                                   nearVal=0.1,
                                                   farVal=100.0)
        _, _, px, _, _ = p.getCameraImage(width=960, height=720,
                                          view_Matrix=view_matrix,
                                          projectionMatrix=proj_matrix,
                                          render=p.ER_BULLET_HARDWARE_OPENGL)
        rgb_array = np.array(px, dtype=np.uint8)
        rgb_array = np.reshape(rgb_array, (720, 960, 4))
        rgb_array = rgb_array[:, :, :3]
        return rgb_array

    def close(self):
        p.disconnect()