import pybullet as p
import pybullet_data
import gym
from gym import spaces
from gym.utils import seeding
import numpy as np
from math import sqrt
import random
import time
import math
import cv2
import torch
import os


def random_crop(imgs, out):
    """
        args:
        imgs: shape (B,C,H,W)
        out: output size (e.g. 84)
    """
    n, c, h, w = imgs.shape
    crop_max = h - out + 1
    w1 = np.random.randint(0, crop_max, n)
    h1 = np.random.randint(0, crop_max, n)
    cropped = np.empty((n, c, out, out), dtype=imgs.dtype)
    for i, (img, w11, h11) in enumerate(zip(imgs, w1, h1)):
        cropped[i] = img[:, h11:h11 + out, w11:w11 + out]
    return cropped


class KukaReachVisualEnv(gym.Env):
    metadata = {
        'render.modes': ['human', 'rgb_array'],
        'video.frames_per_second': 50
    }

    kMaxEpisodeSteps = 700
    kImageSize = {'width': 96, 'height': 96}
    kFinalImageSize = {'width': 84, 'height': 84}

    def __init__(self, is_render=False, is_good_view=False):

        self.is_render = is_render
        self.is_good_view = is_good_view

        if self.is_render:
            p.connect(p.GUI)
        else:
            p.connect(p.DIRECT)

        self.x_low_obs = 0.2
        self.x_high_obs = 0.7
        self.y_low_obs = -0.3
        self.y_high_obs = 0.3
        self.z_low_obs = 0
        self.z_high_obs = 0.55

        self.x_low_action = -0.4
        self.x_high_action = 0.4
        self.y_low_action = -0.4
        self.y_high_action = 0.4
        self.z_low_action = -0.6
        self.z_high_action = 0.3

        self.step_counter = 0

        self.urdf_root_path = pybullet_data.getDataPath()
        # lower limits for null space
        self.lower_limits = [-.967, -2, -2.96, 0.19, -2.96, -2.09, -3.05]
        # upper limits for null space
        self.upper_limits = [.967, 2, 2.96, 2.29, 2.96, 2.09, 3.05]
        # joint ranges for null space
        self.joint_ranges = [5.8, 4, 5.8, 4, 5.8, 4, 6]
        # restposes for null space
        self.rest_poses = [0, 0, 0, 0.5 * math.pi, 0, -math.pi * 0.5 * 0.66, 0]
        # joint damping coefficents
        self.joint_damping = [
            0.00001, 0.00001, 0.00001, 0.00001, 0.00001, 0.00001, 0.00001
        ]

        self.init_joint_positions = [
            0.006418, 0.413184, -0.011401, -1.589317, 0.005379, 1.137684,
            -0.006539
        ]

        self.orientation = p.getQuaternionFromEuler(
            [0., -math.pi, math.pi / 2.])

        self.camera_parameters = {
            'width': 960.,
            'height': 720,
            'fov': 60,
            'near': 0.1,
            'far': 100.,
            'eye_position': [0.59, 0, 0.8],
            'target_position': [0.55, 0, 0.05],
            'camera_up_vector':
                [1, 0, 0],  # I really do not know the parameter's effect.
            'light_direction': [
                0.5, 0, 1
            ],  # the direction is from the light source position to the origin of the world frame.
        }

        self.view_matrix = p.computeViewMatrixFromYawPitchRoll(
            cameraTargetPosition=[0.55, 0, 0.05],
            distance=.7,
            yaw=90,
            pitch=-70,
            roll=0,
            upAxisIndex=2)

        self.projection_matrix = p.computeProjectionMatrixFOV(
            fov=self.camera_parameters['fov'],
            aspect=self.camera_parameters['width'] /
                   self.camera_parameters['height'],
            nearVal=self.camera_parameters['near'],
            farVal=self.camera_parameters['far'])


        p.configureDebugVisualizer(lightPosition=[5, 0, 5])
        p.resetDebugVisualizerCamera(cameraDistance=1.5,
                                     cameraYaw=0,
                                     cameraPitch=-40,
                                     cameraTargetPosition=[0.55, -0.35, 0.2])

        self.action_space = spaces.Box(low=np.array(
            [self.x_low_action, self.y_low_action, self.z_low_action]),
            high=np.array([
                self.x_high_action,
                self.y_high_action,
                self.z_high_action
            ]),
            dtype=np.float32)

        self.observation_space = spaces.Box(low=0, high=1,
                                            shape=(1, self.kFinalImageSize['width'], self.kFinalImageSize['height']))


        self.seed()
        self.reset()

    def seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]

    def reset(self):
        self.step_counter = 0

        p.resetSimulation()
        # p.configureDebugVisualizer(p.COV_ENABLE_RENDERING, 0)
        self.terminated = False
        p.setGravity(0, 0, -10)

        # 这些是周围那些白线，用来观察是否超过了obs的边界
        p.addUserDebugLine(
            lineFromXYZ=[self.x_low_obs, self.y_low_obs, 0],
            lineToXYZ=[self.x_low_obs, self.y_low_obs, self.z_high_obs])
        p.addUserDebugLine(
            lineFromXYZ=[self.x_low_obs, self.y_high_obs, 0],
            lineToXYZ=[self.x_low_obs, self.y_high_obs, self.z_high_obs])
        p.addUserDebugLine(
            lineFromXYZ=[self.x_high_obs, self.y_low_obs, 0],
            lineToXYZ=[self.x_high_obs, self.y_low_obs, self.z_high_obs])
        p.addUserDebugLine(
            lineFromXYZ=[self.x_high_obs, self.y_high_obs, 0],
            lineToXYZ=[self.x_high_obs, self.y_high_obs, self.z_high_obs])

        p.addUserDebugLine(
            lineFromXYZ=[self.x_low_obs, self.y_low_obs, self.z_high_obs],
            lineToXYZ=[self.x_high_obs, self.y_low_obs, self.z_high_obs])
        p.addUserDebugLine(
            lineFromXYZ=[self.x_low_obs, self.y_high_obs, self.z_high_obs],
            lineToXYZ=[self.x_high_obs, self.y_high_obs, self.z_high_obs])
        p.addUserDebugLine(
            lineFromXYZ=[self.x_low_obs, self.y_low_obs, self.z_high_obs],
            lineToXYZ=[self.x_low_obs, self.y_high_obs, self.z_high_obs])
        p.addUserDebugLine(
            lineFromXYZ=[self.x_high_obs, self.y_low_obs, self.z_high_obs],
            lineToXYZ=[self.x_high_obs, self.y_high_obs, self.z_high_obs])

        p.loadURDF(os.path.join(self.urdf_root_path, "plane.urdf"),
                   basePosition=[0, 0, -0.65])
        self.kuka_id = p.loadURDF(os.path.join(self.urdf_root_path,
                                               "kuka_iiwa/model.urdf"),
                                  useFixedBase=True)
        table_uid = p.loadURDF(os.path.join(self.urdf_root_path,
                                            "table/table.urdf"),
                               basePosition=[0.5, 0, -0.65])
        p.changeVisualShape(table_uid, -1, rgbaColor=[1, 1, 1, 1])
        self.object_id = p.loadURDF(os.path.join(self.urdf_root_path,
                                                 "random_urdfs/000/000.urdf"),
                                    basePosition=[
                                        random.uniform(self.x_low_obs,
                                                       self.x_high_obs),
                                        random.uniform(self.y_low_obs,
                                                       self.y_high_obs), 0.01
                                    ])

        self.num_joints = p.getNumJoints(self.kuka_id)

        for i in range(self.num_joints):
            p.resetJointState(
                bodyUniqueId=self.kuka_id,
                jointIndex=i,
                targetValue=self.init_joint_positions[i],
            )

        self.robot_pos_obs = p.getLinkState(self.kuka_id,
                                            self.num_joints - 1)[4]

        p.stepSimulation()


        (_, _, px, _,
         _) = p.getCameraImage(width=960,
                               height=960,
                               viewMatrix=self.view_matrix,
                               projectionMatrix=self.projection_matrix,
                               renderer=p.ER_BULLET_HARDWARE_OPENGL)
        self.images = px

        p.enableJointForceTorqueSensor(bodyUniqueId=self.kuka_id,
                                       jointIndex=self.num_joints - 1,
                                       enableSensor=True)


        self.object_pos = p.getBasePositionAndOrientation(self.object_id)[0]

        self.images = self.images[:, :, :
                                        3]  # the 4th channel is alpha channel, we do not need it.


        return self._process_image(self.images)

    def _process_image(self, image):
        """Convert the RGB pic to gray pic and add a channel 1

        Args:
            image ([type]): [description]
        """

        if image is not None:
            image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
            image = cv2.resize(image, (self.kImageSize['width'], self.kImageSize['height']))[None, :, :] / 255.
            return image
        else:
            return np.zeros((1, self.kImageSize['width'], self.kImageSize['height']))


    def step(self, action):
        dv = 0.005
        dx = action[0] * dv
        dy = action[1] * dv
        dz = action[2] * dv

        self.current_pos = p.getLinkState(self.kuka_id, self.num_joints - 1)[4]
        self.new_robot_pos = [
            self.current_pos[0] + dx, self.current_pos[1] + dy,
            self.current_pos[2] + dz
        ]
        self.robot_joint_positions = p.calculateInverseKinematics(
            bodyUniqueId=self.kuka_id,
            endEffectorLinkIndex=self.num_joints - 1,
            targetPosition=[
                self.new_robot_pos[0], self.new_robot_pos[1],
                self.new_robot_pos[2]
            ],
            targetOrientation=self.orientation,
            jointDamping=self.joint_damping,
        )
        for i in range(self.num_joints):
            p.resetJointState(
                bodyUniqueId=self.kuka_id,
                jointIndex=i,
                targetValue=self.robot_joint_positions[i],
            )
        p.stepSimulation()

        # 在代码开始部分，如果定义了is_good_view，那么机械臂的动作会变慢，方便观察
        if self.is_good_view:
            time.sleep(0.05)

        self.step_counter += 1

        return self._reward()

    def _reward(self):

        # 一定注意是取第4个值，请参考pybullet手册的这个函数返回值的说明
        self.robot_state = p.getLinkState(self.kuka_id, self.num_joints - 1)[4]

        self.object_state = np.array(
            p.getBasePositionAndOrientation(self.object_id)[0]).astype(
            np.float32)

        square_dx = (self.robot_state[0] - self.object_state[0]) ** 2
        square_dy = (self.robot_state[1] - self.object_state[1]) ** 2
        square_dz = (self.robot_state[2] - self.object_state[2]) ** 2

        # 用机械臂末端和物体的距离作为奖励函数的依据
        self.distance = sqrt(square_dx + square_dy + square_dz)
        # print(self.distance)

        x = self.robot_state[0]
        y = self.robot_state[1]
        z = self.robot_state[2]

        # 如果机械比末端超过了obs的空间，也视为done，而且会给予一定的惩罚
        terminated = bool(x < self.x_low_obs or x > self.x_high_obs
                          or y < self.y_low_obs or y > self.y_high_obs
                          or z < self.z_low_obs or z > self.z_high_obs)

        if terminated:
            reward = -0.1
            self.terminated = True

        # 如果机械臂一直无所事事，在最大步数还不能接触到物体，也需要给一定的惩罚
        elif self.step_counter > self.kMaxEpisodeSteps:
            reward = -0.1
            self.terminated = True

        elif self.distance < 0.1:
            reward = 1
            self.terminated = True
        else:
            reward = 0
            self.terminated = False

        info = {'distance:', self.distance}
        (_, _, px, _,
         _) = p.getCameraImage(width=960,
                               height=960,
                               viewMatrix=self.view_matrix,
                               projectionMatrix=self.projection_matrix,
                               renderer=p.ER_BULLET_HARDWARE_OPENGL)
        self.images = px
        self.processed_image = self._process_image(self.images)
        # self.observation=self.robot_state
        self.observation = self.object_state
        return self.processed_image, reward, self.terminated, info

    def close(self):
        p.disconnect()


    def _get_force_sensor_value(self):
        force_sensor_value = p.getJointState(bodyUniqueId=self.kuka_id,
                                             jointIndex=self.num_joints -
                                                        1)[2][2]
        # the first 2 stands for jointReactionForces, the second 2 stands for Fz,
        # the pybullet methods' return is a tuple,so can not
        # index it with str like dict. I think it can be improved
        # that return value is a dict rather than tuple.
        return force_sensor_value


class CustomSkipFrame(gym.Wrapper):
    """ Make a 4 frame skip, so the observation space will change to (4,84,84) from (1,84,84)

    Args:
        gym ([type]): [description]
    """

    def __init__(self, env, skip=4):
        super(CustomSkipFrame, self).__init__(env)
        self.observation_space = spaces.Box(low=0,
                                            high=1,
                                            shape=(skip, self.kFinalImageSize['width'], self.kFinalImageSize['height']))
        self.skip = skip


    def step(self, action):
        total_reward = 0
        states = []
        state, reward, done, info = self.env.step(action)
        for i in range(self.skip):
            if not done:
                state, reward, done, info = self.env.step(action)
                total_reward += reward
                states.append(state)
            else:
                states.append(state)
        states = np.concatenate(states, 0)[None, :, :, :]
        return random_crop(states.astype(np.float32), self.kFinalImageSize['width']), reward, done, info

    def reset(self):
        state = self.env.reset()
        states = np.concatenate([state for _ in range(self.skip)],
                                0)[None, :, :, :]
        return random_crop(states.astype(np.float32), self.kFinalImageSize['width'])


if __name__ == '__main__':
    # 这一部分是做baseline，即让机械臂随机选择动作，看看能够得到的分数
    import matplotlib.pyplot as plt

    env = KukaReachVisualEnv(is_render=False)
    env = CustomSkipFrame(env)
    print(env.observation_space.shape)
    print(env.action_space.shape)
    print(env.action_space.n)
    # for _ in range(20):
    #     action=env.action_space.sample()
    #     print(action)
    #     env.step(action)
    #
    # state = env.reset()
    # print(state.shape)
    # img = state[0][0]
    # plt.imshow(img, cmap='gray')
    # plt.show()



[(24, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (25, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (26, 'numpy.empty', 'np.empty', True, 'import numpy as np\n'), (68, 'pybullet_data.getDataPath', 'pybullet_data.getDataPath', False, 'import pybullet_data\n'), (87, 'pybullet.getQuaternionFromEuler', 'p.getQuaternionFromEuler', True, 'import pybullet as p\n'), (105, 'pybullet.computeViewMatrixFromYawPitchRoll', 'p.computeViewMatrixFromYawPitchRoll', True, 'import pybullet as p\n'), (113, 'pybullet.computeProjectionMatrixFOV', 'p.computeProjectionMatrixFOV', True, 'import pybullet as p\n'), (121, 'pybullet.configureDebugVisualizer', 'p.configureDebugVisualizer', True, 'import pybullet as p\n'), (122, 'pybullet.resetDebugVisualizerCamera', 'p.resetDebugVisualizerCamera', True, 'import pybullet as p\n'), (136, 'gym.spaces.Box', 'spaces.Box', False, 'from gym import spaces\n'), (144, 'gym.utils.seeding.np_random', 'seeding.np_random', False, 'from gym.utils import seeding\n'), (150, 'pybullet.resetSimulation', 'p.resetSimulation', True, 'import pybullet as p\n'), (153, 'pybullet.setGravity', 'p.setGravity', True, 'import pybullet as p\n'), (156, 'pybullet.addUserDebugLine', 'p.addUserDebugLine', True, 'import pybullet as p\n'), (159, 'pybullet.addUserDebugLine', 'p.addUserDebugLine', True, 'import pybullet as p\n'), (162, 'pybullet.addUserDebugLine', 'p.addUserDebugLine', True, 'import pybullet as p\n'), (165, 'pybullet.addUserDebugLine', 'p.addUserDebugLine', True, 'import pybullet as p\n'), (169, 'pybullet.addUserDebugLine', 'p.addUserDebugLine', True, 'import pybullet as p\n'), (172, 'pybullet.addUserDebugLine', 'p.addUserDebugLine', True, 'import pybullet as p\n'), (175, 'pybullet.addUserDebugLine', 'p.addUserDebugLine', True, 'import pybullet as p\n'), (178, 'pybullet.addUserDebugLine', 'p.addUserDebugLine', True, 'import pybullet as p\n'), (190, 'pybullet.changeVisualShape', 'p.changeVisualShape', True, 'import pybullet as p\n'), (200, 'pybullet.getNumJoints', 'p.getNumJoints', True, 'import pybullet as p\n'), (212, 'pybullet.stepSimulation', 'p.stepSimulation', True, 'import pybullet as p\n'), (216, 'pybullet.getCameraImage', 'p.getCameraImage', True, 'import pybullet as p\n'), (223, 'pybullet.enableJointForceTorqueSensor', 'p.enableJointForceTorqueSensor', True, 'import pybullet as p\n'), (262, 'pybullet.calculateInverseKinematics', 'p.calculateInverseKinematics', True, 'import pybullet as p\n'), (278, 'pybullet.stepSimulation', 'p.stepSimulation', True, 'import pybullet as p\n'), (302, 'math.sqrt', 'sqrt', False, 'from math import sqrt\n'), (332, 'pybullet.getCameraImage', 'p.getCameraImage', True, 'import pybullet as p\n'), (344, 'pybullet.disconnect', 'p.disconnect', True, 'import pybullet as p\n'), (367, 'gym.spaces.Box', 'spaces.Box', False, 'from gym import spaces\n'), (48, 'pybullet.connect', 'p.connect', True, 'import pybullet as p\n'), (50, 'pybullet.connect', 'p.connect', True, 'import pybullet as p\n'), (182, 'os.path.join', 'os.path.join', False, 'import os\n'), (184, 'os.path.join', 'os.path.join', False, 'import os\n'), (187, 'os.path.join', 'os.path.join', False, 'import os\n'), (191, 'os.path.join', 'os.path.join', False, 'import os\n'), (203, 'pybullet.resetJointState', 'p.resetJointState', True, 'import pybullet as p\n'), (209, 'pybullet.getLinkState', 'p.getLinkState', True, 'import pybullet as p\n'), (228, 'pybullet.getBasePositionAndOrientation', 'p.getBasePositionAndOrientation', True, 'import pybullet as p\n'), (244, 'cv2.cvtColor', 'cv2.cvtColor', False, 'import cv2\n'), (248, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (257, 'pybullet.getLinkState', 'p.getLinkState', True, 'import pybullet as p\n'), (273, 'pybullet.resetJointState', 'p.resetJointState', True, 'import pybullet as p\n'), (282, 'time.sleep', 'time.sleep', False, 'import time\n'), (291, 'pybullet.getLinkState', 'p.getLinkState', True, 'import pybullet as p\n'), (384, 'numpy.concatenate', 'np.concatenate', True, 'import numpy as np\n'), (127, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (129, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (348, 'pybullet.getJointState', 'p.getJointState', True, 'import pybullet as p\n'), (194, 'random.uniform', 'random.uniform', False, 'import random\n'), (196, 'random.uniform', 'random.uniform', False, 'import random\n'), (245, 'cv2.resize', 'cv2.resize', False, 'import cv2\n'), (294, 'pybullet.getBasePositionAndOrientation', 'p.getBasePositionAndOrientation', True, 'import pybullet as p\n')]