import numpy as np
import gym
from gym import spaces
import time

from manipulator_learning.sim.envs.thing_generic import ThingEnv, get_done_suc_fail


class ThingPickAndPlaceGeneric(ThingEnv):
    def __init__(self,
                 task,
                 camera_in_state,
                 dense_reward,
                 poses_ref_frame,
                 # state_data=('pos', 'obj_pos', 'grip_pos', 'goal_pos'),
                 state_data=('pos', 'obj_pos', 'grip_feedback', 'goal_pos'),
                 max_real_time=5,
                 n_substeps=10,
                 action_multiplier=1.0,
                 reach_radius=.05,
                 gap_between_prev_pos=.2,
                 image_width=160,
                 image_height=120,
                 moving_base=False,
                 t_vel_limit=.3,
                 r_vel_limit=1.5,
                 robot='thing_2_finger',
                 limits_cause_failure=False,
                 failure_causes_done=False,
                 success_causes_done=False,
                 control_frame='w',
                 **kwargs):
        super().__init__(task, robot, camera_in_state,
                         dense_reward, True, poses_ref_frame, state_data, max_real_time=max_real_time,
                         n_substeps=n_substeps, action_multiplier=action_multiplier, gap_between_prev_pos=gap_between_prev_pos,
                         image_width=image_width, image_height=image_height, moving_base=moving_base,
                         failure_causes_done=failure_causes_done, success_causes_done=success_causes_done,
                         control_frame=control_frame, **kwargs)
        self.pos_limits = [[.55, -.45, .64], [1.05, .05, 1.0]]
        self.t_vel_limit = t_vel_limit
        self.r_vel_limit = r_vel_limit
        self.reach_radius = reach_radius
        self.reach_radius_time = .5
        self.reach_radius_start_time = None
        self.in_reach_radius = False
        self.limits_cause_failure = limits_cause_failure
        self.done_success_reward = 100  # hard coded for now, may not work
        self.done_failure_reward = -5   # hard coded for now, may not work

    def _calculate_reward_and_done(self, dense_reward, limit_reached):
        if 'clutter' in self.task:
            return 0, False, False

        if 'lift' in self.task:
            height_above_table_for_suc = self.reach_radius
            table_height = .6247
            des_block_above_table_height = 1.2 * height_above_table_for_suc
            block_pose = self.env._pb_client.getBasePositionAndOrientation(self.env.block_ids[0])
            ee_pose_world = self.env.gripper.manipulator.get_link_pose(
                self.env.gripper.manipulator._tool_link_ind, ref_frame_index=None)
            block_ee_dist = np.linalg.norm(np.array(block_pose[0]) - np.array(ee_pose_world[:3]))
            block_table_dist = block_pose[0][2] - table_height
            block_table_dist_scaled = block_table_dist / des_block_above_table_height
            reward = 3 * min(block_table_dist_scaled, 1.0) + 1 - np.tanh(5.0 * block_ee_dist)
            done_success = False
            if block_table_dist > height_above_table_for_suc:
                if self.reach_radius_start_time is None:
                    self.reach_radius_start_time = self.ep_timesteps
                elif (self.ep_timesteps - self.reach_radius_start_time) * self.real_t_per_ts > self.reach_radius_time:
                    done_success = True
            else:
                self.reach_radius_start_time = None
            done_failure = False
        else:
            block_pose = self.env._pb_client.getBasePositionAndOrientation(self.env.block_ids[0])
            goal_pose = self.env._pb_client.getBasePositionAndOrientation(self.env.goal_id)
            if self.task == 'sort_2' or self.task == 'sort_3':
                block2_pose = self.env._pb_client.getBasePositionAndOrientation(self.env.block_ids[1])
                goal2_pose = self.env._pb_client.getBasePositionAndOrientation(self.env.goal2_id)
            if self.task == 'sort_3':
                block3_pose = self.env._pb_client.getBasePositionAndOrientation(self.env.block_ids[2])
            ee_pose_world = self.env.gripper.manipulator.get_link_pose(
                self.env.gripper.manipulator._tool_link_ind, ref_frame_index=None)
            block_ee_dist = np.linalg.norm(np.array(block_pose[0]) - np.array(ee_pose_world[:3]))
            block_goal_dist = np.linalg.norm(np.array(block_pose[0]) - np.array(goal_pose[0]))
            if self.task == 'sort_2' or self.task == 'sort_3':
                block2_goal_dist = np.linalg.norm(np.array(block2_pose[0]) - np.array(goal2_pose[0]))
            if self.task == 'sort_2':
                reward = 3 * (1 - np.tanh(10.0 * block_goal_dist)) + 1 - np.tanh(10.0 * block2_goal_dist)
            if self.task == 'sort_3':
                block3_goal_dist = np.linalg.norm(np.array(block3_pose[0]) - np.array(goal_pose[0]))
                reward =  (1 - np.tanh(10.0 * block_goal_dist)) + 1 - np.tanh(10.0 * block2_goal_dist) + \
                          1 - np.tanh(10.0 * block3_goal_dist)
            else:
                reward = 3*(1 - np.tanh(10.0 * block_goal_dist)) + 1 - np.tanh(10.0 * block_ee_dist)
            done_success = False
            if (self.task != 'sort_2' and self.task != 'sort_3' and block_goal_dist < self.reach_radius) or \
                (self.task == 'sort_2' and block_goal_dist < self.reach_radius and block2_goal_dist < self.reach_radius) or \
                (self.task == 'sort_3' and block_goal_dist < self.reach_radius and block2_goal_dist < self.reach_radius and
                block3_goal_dist < self.reach_radius):
                if self.reach_radius_start_time is None:
                    self.reach_radius_start_time = self.ep_timesteps
                elif (self.ep_timesteps - self.reach_radius_start_time) * self.real_t_per_ts > self.reach_radius_time:
                    done_success = True
            else:
                self.reach_radius_start_time = None
            done_failure = False
            # num_contact_points = len(self.env._pb_client.getContactPoints(self.env.block_ids[0], self.env.table))

        if self.limits_cause_failure and limit_reached:
            done_failure = True
            done_success = False

        if self.success_causes_done and done_success:
            reward = self.done_success_reward
        if self.failure_causes_done and done_failure:
            reward = self.done_failure_reward

        if dense_reward:
            return reward, done_success, done_failure
        else:
            return done_success, done_success, done_failure

# ----------------------------------------------------------------------------------------------------------
# State Envs
# ----------------------------------------------------------------------------------------------------------

# XY
# ----------------------------------------------------------------------------------------------------------

class ThingPickAndPlaceXYState(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=7, n_substeps=10, dense_reward=True, **kwargs):
        # self.action_space = spaces.Tuple((
        #     spaces.Box(-1, 1, (2,), dtype=np.float32),
        #     spaces.Discrete(2)
        # ))
        self.action_space = spaces.Box(-1, 1, (3,), dtype=np.float32)
        self.observation_space = spaces.Box(-np.inf, np.inf, (7,), dtype=np.float32)
        super().__init__('pick_and_place_xy', False, dense_reward, 'w', max_real_time=max_real_time,
                         n_substeps=n_substeps, **kwargs)

class ThingPickAndPlacePrevPosXYState(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=7, n_substeps=10, dense_reward=True, **kwargs):
        # self.action_space = spaces.Tuple((
        #     spaces.Box(-1, 1, (2,), dtype=np.float32),
        #     spaces.Discrete(2)
        # ))
        self.action_space = spaces.Box(-1, 1, (3,), dtype=np.float32)
        self.observation_space = spaces.Box(-np.inf, np.inf, (16,), dtype=np.float32)
        super().__init__('pick_and_place_xy', False, dense_reward, 'w',
                         state_data=('pos', 'prev_pos', 'obj_pos', 'goal_pos'),
                         max_real_time=max_real_time, n_substeps=n_substeps, **kwargs)

class ThingPickAndPlaceGripPosXYState(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=7, n_substeps=10, dense_reward=True, **kwargs):
        # self.action_space = spaces.Tuple((
        #     spaces.Box(-1, 1, (2,), dtype=np.float32),
        #     spaces.Discrete(2)
        # ))
        self.action_space = spaces.Box(-1, 1, (3,), dtype=np.float32)
        self.observation_space = spaces.Box(-np.inf, np.inf, (8,), dtype=np.float32)
        super().__init__('pick_and_place_xy_vert', False, dense_reward, 'w',
                         state_data=('pos', 'grip_pos', 'obj_pos', 'goal_pos'),
                         max_real_time=max_real_time, n_substeps=n_substeps, **kwargs)

class ThingPickAndPlaceXYImage(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=7, n_substeps=10, dense_reward=True, **kwargs):
        self.action_space = spaces.Tuple((
            spaces.Box(-1, 1, (2,), dtype=np.float32),
            spaces.Discrete(1)
        ))
        self.observation_space = spaces.Dict({
            'obs': spaces.Box(-np.inf, np.inf, (10,), dtype=np.float32),
            'img': spaces.Box(0, 255, (160, 120, 3), dtype=np.uint8),
            'depth': spaces.Box(0, 255, (160, 120), dtype=np.uint8),
        })

        spaces.Dict({"position": spaces.Discrete(2), "velocity": spaces.Discrete(3)})
        super().__init__('pick_and_place_xy', True, dense_reward, 'b', ('prev_pos'), max_real_time=max_real_time,
                         n_substeps=n_substeps, **kwargs)

# XYZ
# ----------------------------------------------------------------------------------------------------------

class ThingPickAndPlaceXYZState(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=8, n_substeps=10, dense_reward=True, **kwargs):
        self.action_space = spaces.Box(-1, 1, (4,), dtype=np.float32)
        self.observation_space = spaces.Box(-np.inf, np.inf, (12,), dtype=np.float32)
        super().__init__('pick_and_place_xyz', False, dense_reward, 'w', max_real_time=max_real_time,
                         n_substeps=n_substeps, **kwargs)
        self.pos_limits = [[.55, -.45, .64], [1.0, .05, 0.8]]

class ThingPickAndPlaceGripPosXYZState(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=8, n_substeps=10, dense_reward=True, **kwargs):
        self.action_space = spaces.Box(-1, 1, (4,), dtype=np.float32)
        self.observation_space = spaces.Box(-np.inf, np.inf, (11,), dtype=np.float32)
        super().__init__('pick_and_place_xyz', False, dense_reward, 'w',
                         state_data=('pos', 'grip_pos', 'obj_pos', 'goal_pos'),
                         max_real_time=max_real_time, n_substeps=n_substeps, **kwargs)
        self.pos_limits = [[.55, -.45, .64], [1.0, .05, 0.8]]

# goal in air, designed to mimic FetchPickAndPlace Env
class ThingPickAndPlaceAirGoalXYZState(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=8, n_substeps=10, dense_reward=True, **kwargs):
        self.action_space = spaces.Box(-1, 1, (4,), dtype=np.float32)
        self.observation_space = spaces.Box(-np.inf, np.inf, (14,), dtype=np.float32)
        super().__init__('pick_and_place_air_xyz', False, dense_reward, 'w',
                         state_data=('pos', 'grip_pos', 'prev_grip_pos', 'obj_pos', 'obj_rot_z', 'goal_pos'),
                         max_real_time=max_real_time, n_substeps=n_substeps, reach_radius=.015, **kwargs)
        self.pos_limits = [[.55, -.45, .66], [0.85, .05, 0.8]]

# 6DOF
# ----------------------------------------------------------------------------------------------------------

class ThingPickAndPlace6DofState(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=8, n_substeps=10, dense_reward=True, **kwargs):
        self.action_space = spaces.Box(-1, 1, (7,), dtype=np.float32)
        self.observation_space = spaces.Box(-np.inf, np.inf, (17,), dtype=np.float32)
        super().__init__('pick_and_place_6dof', False, dense_reward, 'w',
                         state_data=('pos', 'grip_pos', 'obj_pos', 'obj_rot_z', 'goal_pos'),
                         max_real_time=max_real_time, n_substeps=n_substeps, **kwargs)
        self.pos_limits = [[.55, -.45, .64], [.9, .05, 0.8]]

class ThingPickAndPlace6DofLongState(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=8, n_substeps=10, dense_reward=True, **kwargs):
        self.action_space = spaces.Box(-1, 1, (7,), dtype=np.float32)
        self.observation_space = spaces.Box(-np.inf, np.inf, (17,), dtype=np.float32)
        super().__init__('pick_and_place_6dof_long', False, dense_reward, 'w',
                         state_data=('pos', 'grip_pos', 'obj_pos', 'obj_rot_z_sym', 'goal_pos'),
                         max_real_time=max_real_time, n_substeps=n_substeps, reach_radius=.075, **kwargs)
        self.pos_limits = [[.55, -.45, .64], [.9, .05, 0.8]]

class ThingPickAndPlace6DofSmallState(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=8, n_substeps=10, dense_reward=True, **kwargs):
        self.action_space = spaces.Box(-1, 1, (7,), dtype=np.float32)
        self.observation_space = spaces.Box(-np.inf, np.inf, (17,), dtype=np.float32)
        super().__init__('pick_and_place_6dof_small_box', False, dense_reward, 'w',
                         state_data=('pos', 'grip_pos', 'obj_pos', 'obj_rot_z_sym', 'goal_pos'),
                         max_real_time=max_real_time, n_substeps=n_substeps, reach_radius=.075, **kwargs)
        self.pos_limits = [[.55, -.45, .64], [.9, .05, 0.8]]

# goal in air, designed to mimic FetchPickAndPlace Env
class ThingPickAndPlaceAirGoal6DofState(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=8, n_substeps=10, dense_reward=True, **kwargs):
        self.action_space = spaces.Box(-1, 1, (7,), dtype=np.float32)
        self.observation_space = spaces.Box(-np.inf, np.inf, (17,), dtype=np.float32)
        super().__init__('pick_and_place_air_6dof_small_box', False, dense_reward, 'w',
                         state_data=('pos', 'grip_pos', 'obj_pos', 'obj_rot_z_sym', 'goal_pos'),
                         max_real_time=max_real_time, n_substeps=n_substeps, reach_radius=.015, **kwargs)
        self.pos_limits = [[.55, -.45, .64], [1.0, .05, 0.8]]

# ----------------------------------------------------------------------------------------------------------
# SAC-X Envs -- envs designed to mimic environments from SAC-X paper
# ----------------------------------------------------------------------------------------------------------
# Time limit based on appendix of paper --  1200 episodes corresponds to 10 hours = 36000s
# 36000s / 1200 eps = 30s, in their case = 600 timesteps
# objects are on 60cm x 30cm table in any position
# gripper initialized randomly above table-top with height offset 10-20cm above table
# for XYZ envs, gripper always starts facing down
# intention length in paper is set to 150 timesteps = 7.5s
class ThingLiftXYZState(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=18, n_substeps=10, dense_reward=True, action_multiplier=0.1, **kwargs):
        self.action_space = spaces.Box(-1, 1, (4,), dtype=np.float32)
        self.observation_space = spaces.Box(-np.inf, np.inf, (14,), dtype=np.float32)
        # for lift env, reach radius is height above table
        super().__init__('lift_xyz', False, dense_reward, 'w',
                         state_data=('pos', 'grip_pos', 'prev_grip_pos', 'obj_pos', 'obj_rot_z'),
                         max_real_time=max_real_time, n_substeps=n_substeps, 
                         action_multiplier=action_multiplier, reach_radius=.1, robot='thing_2_finger', **kwargs)
        self.pos_limits = [[.55, -.45, .66], [0.85, .05, 0.8]]

class ThingBringXYZState(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=18, n_substeps=10, dense_reward=True, action_multiplier=0.1, **kwargs):
        self.action_space = spaces.Box(-1, 1, (4,), dtype=np.float32)
        self.observation_space = spaces.Box(-np.inf, np.inf, (17,), dtype=np.float32)
        super().__init__('bring_xyz', False, dense_reward, 'w',
                         state_data=('pos', 'grip_pos', 'prev_grip_pos', 'obj_pos', 'obj_rot_z', 'goal_pos'),
                         max_real_time=max_real_time, n_substeps=n_substeps, 
                         action_multiplier=action_multiplier, reach_radius=.015, **kwargs)
        self.pos_limits = [[.55, -.45, .66], [0.85, .05, 0.8]]

# ----------------------------------------------------------------------------------------------------------
# Multiview Paper Envs -- envs to be used with Multiview BC paper
# ----------------------------------------------------------------------------------------------------------

# Use this one to get an expert, and then create an image-based version with shorter max real time
# to use the expert in to generate non-human expert data
class ThingLiftXYZStateMultiview(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=18, n_substeps=10, dense_reward=True, action_multiplier=0.05, **kwargs):
        self.action_space = spaces.Box(-1, 1, (4,), dtype=np.float32)
        self.observation_space = spaces.Box(-np.inf, np.inf, (14,), dtype=np.float32)
        # for lift env, reach radius is height above table
        super().__init__('lift_xyz', False, dense_reward, 'b',
                         state_data=('pos', 'grip_pos', 'prev_grip_pos', 'obj_pos', 'obj_rot_z'),
                         max_real_time=max_real_time, n_substeps=n_substeps,
                         action_multiplier=action_multiplier, reach_radius=.1, robot='thing_2_finger',
                         image_width=64, image_height=48, moving_base=True, control_frame='b', **kwargs)
        self.pos_limits = [[.55, -.45, .66], [0.85, .05, 0.8]]

class ThingLiftXYZState(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=18, n_substeps=10, dense_reward=True, action_multiplier=0.05, **kwargs):
        self.action_space = spaces.Box(-1, 1, (4,), dtype=np.float32)
        self.observation_space = spaces.Box(-np.inf, np.inf, (14,), dtype=np.float32)
        # for lift env, reach radius is height above table
        super().__init__('lift_xyz', False, dense_reward, 'b',
                         state_data=('pos', 'grip_pos', 'prev_grip_pos', 'obj_pos', 'obj_rot_z'),
                         max_real_time=max_real_time, n_substeps=n_substeps,
                         action_multiplier=action_multiplier, reach_radius=.1, robot='thing_2_finger',
                         image_width=64, image_height=48, moving_base=False, control_frame='b', **kwargs)
        self.pos_limits = [[.55, -.45, .66], [0.85, .05, 0.8]]

# todo might have to make this env not have gripper randomizing initial pos
class ThingLiftXYZMultiview(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=6, n_substeps=10, dense_reward=False, action_multiplier=0.05,
                 image_width=64, image_height=48, **kwargs):
        self.action_space = spaces.Box(-1, 1, (4,), dtype=np.float32)
        self.observation_space = spaces.Dict({
            'obs': spaces.Box(-np.inf, np.inf, (9,), dtype=np.float32),
            'img': spaces.Box(0, 255, (image_height, image_width, 3), dtype=np.uint8),
            'depth': spaces.Box(0, 1, (image_height, image_width), dtype=np.float32)
        })
        # for lift env, reach radius is height above table
        super().__init__('lift_xyz', True, dense_reward, 'w',
                         state_data=('pos', 'grip_pos', 'prev_grip_pos'),
                         max_real_time=max_real_time, n_substeps=n_substeps,
                         action_multiplier=action_multiplier, reach_radius=.1, robot='thing_2_finger',
                         image_width=image_width, image_height=image_height, moving_base=True, **kwargs)
        self.pos_limits = [[.55, -.45, .66], [0.85, .05, 0.8]]

class ThingLiftXYZImage(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=6, n_substeps=10, dense_reward=False, action_multiplier=0.05,
                 image_width=64, image_height=48, **kwargs):
        self.action_space = spaces.Box(-1, 1, (4,), dtype=np.float32)
        self.observation_space = spaces.Dict({
            'obs': spaces.Box(-np.inf, np.inf, (9,), dtype=np.float32),
            'img': spaces.Box(0, 255, (image_height, image_width, 3), dtype=np.uint8),
            'depth': spaces.Box(0, 1, (image_height, image_width), dtype=np.float32)
        })
        # for lift env, reach radius is height above table
        super().__init__('lift_xyz', True, dense_reward, 'w',
                         state_data=('pos', 'grip_pos', 'prev_grip_pos'),
                         max_real_time=max_real_time, n_substeps=n_substeps,
                         action_multiplier=action_multiplier, reach_radius=.1, robot='thing_2_finger',
                         image_width=image_width, image_height=image_height, moving_base=False, **kwargs)
        self.pos_limits = [[.55, -.45, .66], [0.85, .05, 0.8]]

# ----------------------------------------------------------------------------------------------------------
# Image Envs
# ----------------------------------------------------------------------------------------------------------

class ThingPickAndPlace6DofSmallImage(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=8, n_substeps=10, dense_reward=True,
                 image_width=64, image_height=48, **kwargs):
        self.action_space = spaces.Box(-1, 1, (7,), dtype=np.float32)
        self.observation_space = spaces.Dict({
            'obs': spaces.Box(-np.inf, np.inf, (9,), dtype=np.float32),
            'img': spaces.Box(0, 255, (image_height, image_width, 3), dtype=np.uint8),
            'depth': spaces.Box(0, 1, (image_height, image_width), dtype=np.float32)
        })
        super().__init__('pick_and_place_6dof_small_box_image', True, dense_reward, 'b',
                         state_data=('pos', 'grip_pos'),
                         max_real_time=max_real_time, n_substeps=n_substeps, reach_radius=.075,
                         image_width=image_width, image_height=image_height, control_frame='b', **kwargs)
        self.pos_limits = [[.55, -.45, .64], [.9, .05, 0.8]]

class ThingPickAndPlace6DofSmall160120Image(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=8, n_substeps=10, dense_reward=True, **kwargs):
        image_height = 120
        image_width = 160
        self.action_space = spaces.Box(-1, 1, (7,), dtype=np.float32)
        self.observation_space = spaces.Dict({
            'obs': spaces.Box(-np.inf, np.inf, (9,), dtype=np.float32),
            'img': spaces.Box(0, 255, (image_height, image_width, 3), dtype=np.uint8),
            'depth': spaces.Box(0, 1, (image_height, image_width), dtype=np.float32)
        })
        super().__init__('pick_and_place_6dof_small_box_image', True, dense_reward, 'b',
                         state_data=('pos', 'grip_pos'),
                         max_real_time=max_real_time, n_substeps=n_substeps, reach_radius=.075,
                         image_width=image_width, image_height=image_height, **kwargs)
        self.pos_limits = [[.55, -.45, .64], [.9, .05, 0.8]]

class ThingPickAndPlace6DofSmallMultiview(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=8, n_substeps=10, dense_reward=True,
                 image_width=64, image_height=48, **kwargs):
        self.action_space = spaces.Box(-1, 1, (7,), dtype=np.float32)
        self.observation_space = spaces.Dict({
            'obs': spaces.Box(-np.inf, np.inf, (9,), dtype=np.float32),
            'img': spaces.Box(0, 255, (image_height, image_width, 3), dtype=np.uint8),
            'depth': spaces.Box(0, 1, (image_height, image_width), dtype=np.float32)
        })
        super().__init__('pick_and_place_6dof_small_box_image', True, dense_reward, 'b',
                         state_data=('pos', 'grip_pos'),
                         max_real_time=max_real_time, n_substeps=n_substeps, reach_radius=.075,
                         image_width=image_width, image_height=image_height, moving_base=True,
                         r_vel_limit=1.5, **kwargs)
        self.pos_limits = [[.55, -.45, .64], [.9, .05, 0.8]]

class ThingSort2Multiview(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=12, n_substeps=10, dense_reward=True,
                 image_width=64, image_height=48, **kwargs):
        self.action_space = spaces.Box(-1, 1, (7,), dtype=np.float32)
        self.observation_space = spaces.Dict({
            'obs': spaces.Box(-np.inf, np.inf, (9,), dtype=np.float32),
            'img': spaces.Box(0, 255, (image_height, image_width, 3), dtype=np.uint8),
            'depth': spaces.Box(0, 1, (image_height, image_width), dtype=np.float32)
        })
        super().__init__('sort_2', True, dense_reward, 'b',
                         state_data=('pos', 'grip_pos'),
                         max_real_time=max_real_time, n_substeps=n_substeps, reach_radius=.075,
                         image_width=image_width, image_height=image_height, moving_base=True,
                         r_vel_limit=1.5, **kwargs)
        self.pos_limits = [[.55, -.45, .64], [1.0, .05, 0.8]]

class ThingSort3Multiview(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=15, n_substeps=10, dense_reward=True,
                 image_width=64, image_height=48, **kwargs):
        self.action_space = spaces.Box(-1, 1, (7,), dtype=np.float32)
        self.observation_space = spaces.Dict({
            'obs': spaces.Box(-np.inf, np.inf, (9,), dtype=np.float32),
            'img': spaces.Box(0, 255, (image_height, image_width, 3), dtype=np.uint8),
            'depth': spaces.Box(0, 1, (image_height, image_width), dtype=np.float32)
        })
        super().__init__('sort_3', True, dense_reward, 'b',
                         state_data=('pos', 'grip_pos'),
                         max_real_time=max_real_time, n_substeps=n_substeps, reach_radius=.1,
                         image_width=image_width, image_height=image_height, moving_base=True,
                         r_vel_limit=1.5, **kwargs)
        self.pos_limits = [[.55, -.45, .64], [1.0, .05, 0.8]]

class ThingPickAndPlaceClutterMultiview(ThingPickAndPlaceGeneric):
    def __init__(self, max_real_time=12, n_substeps=10, dense_reward=True,
                 image_width=64, image_height=48, **kwargs):
        self.action_space = spaces.Box(-1, 1, (7,), dtype=np.float32)
        self.observation_space = spaces.Dict({
            'obs': spaces.Box(-np.inf, np.inf, (9,), dtype=np.float32),
            'img': spaces.Box(0, 255, (image_height, image_width, 3), dtype=np.uint8),
            'depth': spaces.Box(0, 1, (image_height, image_width), dtype=np.float32)
        })
        super().__init__('pick_and_place_clutter', True, dense_reward, 'b',
                         state_data=('pos', 'grip_pos'),
                         max_real_time=max_real_time, n_substeps=n_substeps, reach_radius=.075,
                         image_width=image_width, image_height=image_height, moving_base=True,
                         r_vel_limit=1.5, **kwargs)
        self.pos_limits = [[.55, -.45, .64], [1.0, .05, 0.8]]
