"""Base Task class."""

import collections
import os
import random
import string
import tempfile

import cv2
import numpy as np
from cliport.tasks import cameras
from cliport.tasks import primitives
from cliport.tasks.grippers import Suction
from cliport.utils import utils
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
import pybullet as p
from typing import Tuple, List
import re

class Task():
    """Base Task class."""

    def __init__(self):
        self.ee = Suction
        self.mode = 'train'
        self.sixdof = False
        self.primitive = primitives.PickPlace()
        self.oracle_cams = cameras.Oracle.CONFIG

        # Evaluation epsilons (for pose evaluation metric).
        self.pos_eps = 0.01
        self.rot_eps = np.deg2rad(15)

        # for piles
        self.num_blocks = 50

        # Workspace bounds.
        self.pix_size = 0.003125
        self.bounds = np.array([[0.25, 0.75], [-0.5, 0.5], [0, 0.3]])
        self.zone_bounds = np.copy(self.bounds)

        self.goals = []
        self.lang_goals = []
        self.obj_points_cache = {}

        self.task_completed_desc = "task completed."
        self.progress = 0
        self._rewards = 0

        self.train_set = np.arange(0, 14)
        self.test_set = np.arange(14, 20)
        self.assets_root = None
        self.homogeneous = False

    def reset(self, env):
        if not self.assets_root:
            raise ValueError('assets_root must be set for task, '
                             'call set_assets_root().')
        self.goals = []
        self.lang_goals = []
        self.progress = 0  # Task progression metric in range [0, 1].
        self._rewards = 0  # Cumulative returned rewards.
        self.obj_points_cache = {}

    def additional_reset(self):
        # Additional changes to make the environment adaptable
        if 'bowl' in self.lang_template:
            # IMPORTANT: increase position tolerance for bowl placement
            self.pos_eps = 0.05

        if 'piles' in self.lang_template:
            # IMPORTANT: Define the primitive to be push and ee to be spatula for tasks involving piles
            self.ee = Spatula
            self.primitive = primitives.push

        if 'rope' in self.lang_template:
            self.primitive = primitives.PickPlace(height=0.02, speed=0.001)
            self.pos_eps = 0.02

    # -------------------------------------------------------------------------
    # Oracle Agent
    # -------------------------------------------------------------------------

    def oracle(self, env):
        """Oracle agent."""
        OracleAgent = collections.namedtuple('OracleAgent', ['act'])

        def act(obs, info):
            """Calculate action."""

            # Oracle uses perfect RGB-D orthographic images and segmentation masks.
            _, hmap, obj_mask = self.get_true_image(env)

            # Unpack next goal step.
            objs, matches, targs, replace, rotations, _, _, _ = self.goals[0]

            for j, targ in enumerate(targs):
                # add default orientation if missing
                if len(targ) == 3 and (type(targs[j][0]) is float or type(targs[j][0]) is np.float32):
                    targs[j] = (targs[j], (0,0,0,1))

            # Match objects to targets without replacement.
            if not replace:

                # Modify a copy of the match matrix.
                matches = matches.copy()

                # Ignore already matched objects.
                for i in range(len(objs)):
                    if type(objs[i]) is int:
                        objs[i] = (objs[i], (False, None))

                    object_id, (symmetry, _) = objs[i]
                    pose = p.getBasePositionAndOrientation(object_id)
                    targets_i = np.argwhere(matches[i, :]).reshape(-1)
                    for j in targets_i:
                        if self.is_match(pose, targs[j], symmetry):
                            matches[i, :] = 0
                            matches[:, j] = 0

            # Get objects to be picked (prioritize farthest from nearest neighbor).
            nn_dists = []
            nn_targets = []
            for i in range(len(objs)):
                if type(objs[i]) is int:
                    objs[i] = (objs[i], (False, None))

                object_id, (symmetry, _) = objs[i]
                xyz, _ = p.getBasePositionAndOrientation(object_id)
                targets_i = np.argwhere(matches[i, :]).reshape(-1)
                if len(targets_i) > 0:

                    targets_xyz = np.float32([targs[j][0] for j in targets_i])
                    dists = np.linalg.norm(
                        targets_xyz - np.float32(xyz).reshape(1, 3), axis=1)
                    nn = np.argmin(dists)
                    nn_dists.append(dists[nn])
                    nn_targets.append(targets_i[nn])

                # Handle ignored objects.
                else:
                    nn_dists.append(0)
                    nn_targets.append(-1)
            order = np.argsort(nn_dists)[::-1]

            # Filter out matched objects.
            order = [i for i in order if nn_dists[i] > 0]

            pick_mask = None
            for pick_i in order:
                pick_mask = np.uint8(obj_mask == objs[pick_i][0])

                # Erode to avoid picking on edges.
                pick_mask = cv2.erode(pick_mask, np.ones((3, 3), np.uint8))

                if np.sum(pick_mask) > 0:
                    break

            # Trigger task reset if no object is visible.
            if pick_mask is None or np.sum(pick_mask) == 0:
                self.goals = []
                self.lang_goals = []
                print('Object for pick is not visible. Skipping demonstration.')
                return

            # Get picking pose.
            pick_prob = np.float32(pick_mask)
            pick_pix = utils.sample_distribution(pick_prob)
            # For "deterministic" demonstrations on insertion-easy, use this:
            pick_pos = utils.pix_to_xyz(pick_pix, hmap,
                                        self.bounds, self.pix_size)
            pick_pose = (np.asarray(pick_pos), np.asarray((0, 0, 0, 1)))

            # Get placing pose.
            targ_pose = targs[nn_targets[pick_i]]
            obj_pose = p.getBasePositionAndOrientation(objs[pick_i][0])
            if not self.sixdof:
                obj_euler = utils.quatXYZW_to_eulerXYZ(obj_pose[1])
                obj_quat = utils.eulerXYZ_to_quatXYZW((0, 0, obj_euler[2]))
                obj_pose = (obj_pose[0], obj_quat)
            world_to_pick = utils.invert(pick_pose)
            obj_to_pick = utils.multiply(world_to_pick, obj_pose)
            pick_to_obj = utils.invert(obj_to_pick)

            if len(targ_pose) == 3 and (type(targ_pose[0]) is float or type(targ_pose[0]) is np.float32):
                # add default orientation if missing
                targ_pose = (targ_pose, (0,0,0,1))

            place_pose = utils.multiply(targ_pose, pick_to_obj)

            # Rotate end effector?
            if not rotations:
                place_pose = (place_pose[0], (0, 0, 0, 1))

            place_pose = (np.asarray(place_pose[0]), np.asarray(place_pose[1]))

            return {'pose0': pick_pose, 'pose1': place_pose}

        return OracleAgent(act)

    # -------------------------------------------------------------------------
    # Reward Function and Task Completion Metrics
    # -------------------------------------------------------------------------

    def reward(self):
        """Get delta rewards for current timestep.

        Returns:
          A tuple consisting of the scalar (delta) reward.
        """
        reward, info = 0, {}

        # Unpack next goal step.
        objs, matches, targs, replace, _, metric, params, max_reward = self.goals[0]

        # Evaluate by matching object poses.
        step_reward = 0

        if metric == 'pose':
            for i in range(len(objs)):
                object_id, (symmetry, _) = objs[i]
                pose = p.getBasePositionAndOrientation(object_id)
                targets_i = np.argwhere(matches[i, :])
                if len(targets_i) > 0:
                    targets_i = targets_i.reshape(-1)
                    for j in targets_i:
                        target_pose = targs[j]
                        if self.is_match(pose, target_pose, symmetry):
                            step_reward += max_reward / len(objs)
                            print(f"object {i} match with target {j} rew: {step_reward:.3f}")
                            break

        # Evaluate by measuring object intersection with zone.
        elif metric == 'zone':
            zone_pts, total_pts = 0, 0
            zones = params

            if len(self.obj_points_cache) == 0 or objs[0][0] not in self.obj_points_cache:
                for obj_id, _ in objs:
                    self.obj_points_cache[obj_id] = self.get_box_object_points(obj_id)

            for zone_idx, (zone_pose, zone_size) in enumerate(zones):
                # Count valid points in zone.
                for (obj_id, _) in objs:
                    pts = self.obj_points_cache[obj_id]
                    obj_pose = p.getBasePositionAndOrientation(obj_id)
                    world_to_zone = utils.invert(zone_pose)
                    obj_to_zone = utils.multiply(world_to_zone, obj_pose)
                    pts = np.float32(utils.apply(obj_to_zone, pts))
                    # if type(zone_size) is int:
                    #      print("closest point:", p.getClosestPoints(obj_id, zone_size, 0.1))
                    #     valid_pts = len(p.getClosestPoints(obj_id, zone_size, 0.1)) > 0

                    if len(zone_size) > 1:
                        valid_pts = np.logical_and.reduce([
                            pts[0, :] > -zone_size[0] / 2, pts[0, :] < zone_size[0] / 2,
                            pts[1, :] > -zone_size[1] / 2, pts[1, :] < zone_size[1] / 2,
                            pts[2, :] < self.zone_bounds[2, 1]])

                    zone_pts += np.sum(np.float32(valid_pts))
                    total_pts += pts.shape[1]

            if total_pts > 0:
                step_reward = max_reward * (zone_pts / total_pts)

        # Get cumulative rewards and return delta.
        reward = self.progress + step_reward - self._rewards
        self._rewards = self.progress + step_reward

        # Move to next goal step if current goal step is complete.
        if np.abs(max_reward - step_reward) < 0.01:
            self.progress += max_reward  # Update task progress.
            self.goals.pop(0)
            if len(self.lang_goals) > 0:
                self.lang_goals.pop(0)

        return reward, info

    def done(self):
        """Check if the task is done or has failed.

        Returns:
          True if the episode should be considered a success.
        """
        return (len(self.goals) == 0) or (self._rewards > 0.99)
        # return zone_done or defs_done or goal_done

    # -------------------------------------------------------------------------
    # Environment Helper Functions
    # -------------------------------------------------------------------------

    def is_match(self, pose0, pose1, symmetry):
        """Check if pose0 and pose1 match within a threshold.
        pose0 and pose1 should both be tuples of (translation, rotation).
        Return true if the pose translation and orientation errors are below certain thresholds"""
        if len(pose1) == 3 and (not hasattr(pose1[0], '__len__')):
            # add default orientation if missing
            pose1 = (pose1, (0,0,0,1))
        # print(len(pose1) == 3, not hasattr(pose1[0], '__len__'))
        # print(pose1, pose0)
        # Get translational error.
        diff_pos = np.float32(pose0[0][:2]) - np.float32(pose1[0][:2])
        dist_pos = np.linalg.norm(diff_pos)

        # Get rotational error around z-axis (account for symmetries).
        diff_rot = 0
        if symmetry > 0:
            rot0 = np.array(utils.quatXYZW_to_eulerXYZ(pose0[1]))[2]
            rot1 = np.array(utils.quatXYZW_to_eulerXYZ(pose1[1]))[2]
            diff_rot = np.abs(rot0 - rot1) % symmetry
            if diff_rot > (symmetry / 2):
                diff_rot = symmetry - diff_rot

        return (dist_pos < self.pos_eps) and (diff_rot < self.rot_eps)

    def get_true_image(self, env):
        """Get RGB-D orthographic heightmaps and segmentation masks."""

        # Capture near-orthographic RGB-D images and segmentation masks.
        color, depth, segm = env.render_camera(self.oracle_cams[0])

        # Combine color with masks for faster processing.
        color = np.concatenate((color, segm[Ellipsis, None]), axis=2)

        # Reconstruct real orthographic projection from point clouds.
        hmaps, cmaps = utils.reconstruct_heightmaps(
            [color], [depth], self.oracle_cams, self.bounds, self.pix_size)

        # Split color back into color and masks.
        cmap = np.uint8(cmaps)[0, Ellipsis, :3]
        hmap = np.float32(hmaps)[0, Ellipsis]
        mask = np.int32(cmaps)[0, Ellipsis, 3:].squeeze()
        return cmap, hmap, mask

    def get_random_pose(self, env, obj_size=0.1, **kwargs) -> (List, List):
        """
        Get random collision-free object pose within workspace bounds.
        :param obj_size: (3, ) contains the object size in x,y,z dimensions
        return: translation (3, ), rotation (4, ) """

        # Get erosion size of object in pixels.
        max_size = np.sqrt(obj_size[0] ** 2 + obj_size[1] ** 2)
        erode_size = int(np.round(max_size / self.pix_size))

        _, hmap, obj_mask = self.get_true_image(env)

        # Randomly sample an object pose within free-space pixels.
        free = np.ones(obj_mask.shape, dtype=np.uint8)
        for obj_ids in env.obj_ids.values():
            for obj_id in obj_ids:
                free[obj_mask == obj_id] = 0
        free[0, :], free[:, 0], free[-1, :], free[:, -1] = 0, 0, 0, 0
        free = cv2.erode(free, np.ones((erode_size, erode_size), np.uint8))

        # if np.sum(free) == 0:
        #     return None, None

        if np.sum(free) == 0:
            # avoid returning None
            pix = (obj_mask.shape[0] // 2, obj_mask.shape[1] // 2)
        else:
            pix = utils.sample_distribution(np.float32(free))
        pos = utils.pix_to_xyz(pix, hmap, self.bounds, self.pix_size)

        if len(obj_size) == 2:
            print("Should have z dimension in obj_size as well.")
            pos = [pos[0], pos[1], 0.05]
        else:
            pos = [pos[0], pos[1], obj_size[2] / 2]
        theta = np.random.rand() * 2 * np.pi
        rot = utils.eulerXYZ_to_quatXYZW((0, 0, theta))
        return pos, rot

    def get_lang_goal(self):
        if len(self.lang_goals) == 0:
            return self.task_completed_desc
        else:
            return self.lang_goals[0]

    def get_reward(self):
        return float(self._rewards)
    
    def add_corner_anchor_for_pose(self, env, pose):
        corner_template = 'corner/corner-template.urdf'
        replace = {'DIMX': (0.04,), 'DIMY': (0.04,)}

        # IMPORTANT: REPLACE THE TEMPLATE URDF
        corner_urdf = self.fill_template(corner_template, replace)
        if len(pose) != 2:
            pose = [pose,(0,0,0,1)]
        env.add_object(corner_urdf, pose, 'fixed')


    def get_target_sample_surface_points(self, model, scale, pose, num_points=50):
        import trimesh
        mesh = trimesh.load_mesh(model)
        points = trimesh.sample.volume_mesh(mesh, num_points * 3)
        points = points[:num_points]
        points = points * np.array(scale)
        points = utils.apply(pose, points.T)
        poses = [((x,y,z),(0,0,0,1)) for x, y, z in zip(points[0], points[1], points[2])]
        return poses
    # -------------------------------------------------------------------------
    # Helper Functions
    # -------------------------------------------------------------------------
    def check_require_obj(self, path):
        return os.path.exists(path.replace(".urdf", ".obj"))

    def fill_template(self, template, replace):
        """Read a file and replace key strings.
        NOTE: This function must be called if a URDF has template in its name """

        full_template_path = os.path.join(self.assets_root, template)
        if not os.path.exists(full_template_path) or (self.check_require_obj(full_template_path) and 'template' not in full_template_path):
            return template

        with open(full_template_path, 'r') as file:
            fdata = file.read()

        for field in replace:
            # if  not hasattr(replace[field], '__len__'):
            #     replace[field] = (replace[field], )

            for i in range(len(replace[field])):
                fdata = fdata.replace(f'{field}{i}', str(replace[field][i]))

            if field == 'COLOR':
                # handle gpt
                pattern = r'<color rgba="(.*?)"/>'
                code_string = re.findall(pattern, fdata)
                if type(replace[field]) is str:
                    replace[field] = utils.COLORS[replace[field]]
                for to_replace_color in  code_string:
                    fdata = fdata.replace(f'{to_replace_color}', " ".join([str(x) for x in list(replace[field]) + [1]]))
            
        alphabet = string.ascii_lowercase + string.digits
        rname = ''.join(random.choices(alphabet, k=16))
        tmpdir = tempfile.gettempdir()
        template_filename = os.path.split(template)[-1]
        fname = os.path.join(tmpdir, f'{template_filename}.{rname}')
        with open(fname, 'w') as file:
            file.write(fdata)
        return fname

    def get_random_size(self, min_x, max_x, min_y, max_y, min_z, max_z) -> Tuple:
        """Get random box size."""
        size = np.random.rand(3)
        size[0] = size[0] * (max_x - min_x) + min_x
        size[1] = size[1] * (max_y - min_y) + min_y
        size[2] = size[2] * (max_z - min_z) + min_z
        return tuple(size)

    def get_box_object_points(self, obj):
        obj_shape = p.getVisualShapeData(obj)
        obj_dim = obj_shape[0][3]
        obj_dim = tuple(d for d in obj_dim)
        xv, yv, zv = np.meshgrid(
            np.arange(-obj_dim[0] / 2, obj_dim[0] / 2, 0.02),
            np.arange(-obj_dim[1] / 2, obj_dim[1] / 2, 0.02),
            np.arange(-obj_dim[2] / 2, obj_dim[2] / 2, 0.02),
            sparse=False, indexing='xy')
        return np.vstack((xv.reshape(1, -1), yv.reshape(1, -1), zv.reshape(1, -1)))

    def get_sphere_object_points(self, obj):
        return self.get_box_object_points(obj)

    def get_mesh_object_points(self, obj):
        mesh = p.getMeshData(obj)
        mesh_points = np.array(mesh[1])
        mesh_dim = np.vstack((mesh_points.min(axis=0), mesh_points.max(axis=0)))
        xv, yv, zv = np.meshgrid(
            np.arange(mesh_dim[0][0], mesh_dim[1][0], 0.02),
            np.arange(mesh_dim[0][1], mesh_dim[1][1], 0.02),
            np.arange(mesh_dim[0][2], mesh_dim[1][2], 0.02),
            sparse=False, indexing='xy')
        return np.vstack((xv.reshape(1, -1), yv.reshape(1, -1), zv.reshape(1, -1)))

    def color_random_brown(self, obj):
        shade = np.random.rand() + 0.5
        color = np.float32([shade * 156, shade * 117, shade * 95, 255]) / 255
        p.changeVisualShape(obj, -1, rgbaColor=color)

    def set_assets_root(self, assets_root):
        self.assets_root = assets_root

    def zip_obj_ids(self, obj_ids, symmetries):
        if type(obj_ids[0]) is tuple:
            return obj_ids

        if  symmetries is None:
             symmetries = [0.] * len(obj_ids)
        objs = []

        for obj_id, symmetry in zip(obj_ids, symmetries):
            objs.append((obj_id, (symmetry, None)))
        return objs

    def add_goal(self, objs, matches, targ_poses, replace, rotations, metric, params, step_max_reward, symmetries=None, **kwargs):
        """ Add the goal to the environment
        - objs (List of Tuple [(obj_id, (float, None))] ): object ID, (the radians that the object is symmetric over, None). Do not pass in `(object id, object pose)` as the wrong tuple. or `object id` (such as `containers[i][0]`).
        - matches (Binary Matrix): a binary matrix that denotes which object is matched with which target. This matrix has dimension len(objs) x len(targ_poses).
        - targ_poses (List of Poses [(translation, rotation)] ): a list of target poses of tuple (translation, rotation). Don't pass in object IDs such as `bowls[i-1][0]` or  `[stands[i][0]]`.
        - replace (Boolean): whether each object can match with one unique target.   This is important if we have one target and multiple objects. If it's set to be false, then any object matching with the target will satisfy.
        - rotations (Boolean): whether the placement action has a rotation degree of freedom.
        - metric (`pose` or `zone`): `pose` or `zone` that the object needs to be transported to. Example: `pose`.
        - params ([(zone_target, zone_size)])): has to be [(zone_target, zone_size)] if the metric is `zone` where obj_pts is a dictionary that maps object ID to points.
        - step_max_reward (float): the maximum reward of matching all the objects with all the target poses.
        """
        objs = self.zip_obj_ids(objs, symmetries)
        self.goals.append((objs, matches, targ_poses, replace, rotations,
                           metric, params, step_max_reward))


    def make_piles(self, env, block_color=None, *args, **kwargs):
        """
        add the piles objects for tasks involving piles
        """
        obj_ids = []
        for _ in range(self.num_blocks):
            rx = self.bounds[0, 0] + 0.15 + np.random.rand() * 0.2
            ry = self.bounds[1, 0] + 0.4 + np.random.rand() * 0.2
            xyz = (rx, ry, 0.01)
            theta = np.random.rand() * 2 * np.pi
            xyzw = utils.eulerXYZ_to_quatXYZW((0, 0, theta))
            obj_id = env.add_object('block/small.urdf', (xyz, xyzw))
            if block_color is not None:
                p.changeVisualShape(obj_id, -1, rgbaColor=block_color + [1])

            obj_ids.append(obj_id)
        return obj_ids

    def make_rope(self, *args, **kwargs):
        return self.make_ropes(*args, **kwargs)

    def make_ropes(self, env, corners, radius=0.005, n_parts = 20, color_name='red', *args, **kwargs):
        """ add cables simulation for tasks that involve cables """
        # Get corner points of square.
        
        # radius = 0.005
        length = 2 * radius * n_parts * np.sqrt(2)
        corner0, corner1 = corners
        # Add cable (series of articulated small blocks).
        increment = (np.float32(corner1) - np.float32(corner0)) / n_parts
        position, _ = self.get_random_pose(env, (0.1, 0.1, 0.1))
        position = np.float32(position)
        part_shape = p.createCollisionShape(p.GEOM_BOX, halfExtents=[radius] * 3)
        part_visual = p.createVisualShape(p.GEOM_SPHERE, radius=radius * 1.5)
        parent_id = -1
        targets = []
        objects = []

        for i in range(n_parts):
            position[2] += np.linalg.norm(increment)
            part_id = p.createMultiBody(0.1, part_shape, part_visual,
                                        basePosition=position)
            if parent_id > -1:
                constraint_id = p.createConstraint(
                    parentBodyUniqueId=parent_id,
                    parentLinkIndex=-1,
                    childBodyUniqueId=part_id,
                    childLinkIndex=-1,
                    jointType=p.JOINT_POINT2POINT,
                    jointAxis=(0, 0, 0),
                    parentFramePosition=(0, 0, np.linalg.norm(increment)),
                    childFramePosition=(0, 0, 0))
                p.changeConstraint(constraint_id, maxForce=100)

            if (i > 0) and (i < n_parts - 1):
                color = utils.COLORS[color_name] + [1]
                p.changeVisualShape(part_id, -1, rgbaColor=color)

            env.obj_ids['rigid'].append(part_id)
            parent_id = part_id
            target_xyz = np.float32(corner0) + i * increment + increment / 2
            objects.append((part_id, (0, None)))
            targets.append((target_xyz, (0, 0, 0, 1)))

            if  hasattr(env, 'record_cfg') and 'blender_render' in env.record_cfg and env.record_cfg['blender_render']:
                sphere_template = os.path.join(self.assets_root, 'sphere/sphere_rope.urdf')
                env.blender_recorder.register_object(part_id, os.path.join(self.assets_root, 'sphere/sphere_rope.urdf'))


        matches = np.clip(np.eye(n_parts) + np.eye(n_parts)[::-1], 0, 1)
        return objects, targets, matches


    def get_kitting_shapes(self, n_objects):
        if self.mode == 'train':
            obj_shapes = np.random.choice(self.train_set, n_objects)
        else:
            if self.homogeneous:
                obj_shapes = [np.random.choice(self.test_set)] * n_objects
            else:
                obj_shapes = np.random.choice(self.test_set, n_objects)

        return obj_shapes


    def make_kitting_objects(self, env, targets, obj_shapes, n_objects, colors):
        symmetry = [
            2 * np.pi, 2 * np.pi, 2 * np.pi / 3, np.pi / 2, np.pi / 2, 2 * np.pi,
            np.pi, 2 * np.pi / 5, np.pi, np.pi / 2, 2 * np.pi / 5, 0, 2 * np.pi,
            2 * np.pi, 2 * np.pi, 2 * np.pi, 0, 2 * np.pi / 6, 2 * np.pi, 2 * np.pi
        ]
        objects = []
        matches = []
        template = 'kitting/object-template.urdf'

        for i in range(n_objects):
            shape = obj_shapes[i]
            size = (0.08, 0.08, 0.02)
            pose = self.get_random_pose(env, size)
            fname = f'{shape:02d}.obj'
            fname = os.path.join(self.assets_root, 'kitting', fname)
            scale = [0.003, 0.003, 0.001]  # .0005
            replace = {'FNAME': (fname,), 'SCALE': scale, 'COLOR': colors[i]}

            # IMPORTANT: REPLACE THE TEMPLATE URDF
            urdf = self.fill_template(template, replace)
            block_id = env.add_object(urdf, pose)
            objects.append((block_id, (symmetry[shape], None)))
            match = np.zeros(len(targets))
            match[np.argwhere(obj_shapes == shape).reshape(-1)] = 1
            matches.append(match)
        return objects, matches

    def spawn_box(self):
        """Palletizing: spawn another box in the workspace if it is empty."""
        workspace_empty = True
        if self.goals:
            for obj in self.goals[0][0]:
                obj_pose = p.getBasePositionAndOrientation(obj[0])
                workspace_empty = workspace_empty and ((obj_pose[0][1] < -0.5) or
                                                       (obj_pose[0][1] > 0))
            if not self.steps:
                self.goals = []
                print('Palletized boxes toppled. Terminating episode.')
                return

            if workspace_empty:
                obj = self.steps[0]
                theta = np.random.random() * 2 * np.pi
                rotation = utils.eulerXYZ_to_quatXYZW((0, 0, theta))
                p.resetBasePositionAndOrientation(obj, [0.5, -0.25, 0.1], rotation)
                self.steps.pop(0)

        # Wait until spawned box settles.
        for _ in range(480):
            p.stepSimulation()

    def get_asset_full_path(self, path):
        return path