import numpy as np

from igibson.external.pybullet_tools.utils import aabb_union, get_aabb, get_all_links
from igibson.object_states.object_state_base import CachingEnabledObjectState


class AABB(CachingEnabledObjectState):
    def _compute_value(self):
        body_id = self.obj.get_body_id()
        all_links = get_all_links(body_id)
        aabbs = [get_aabb(body_id, link=link) for link in all_links]
        aabb_low, aabb_hi = aabb_union(aabbs)

        if not hasattr(self.obj, "category") or self.obj.category != "floors" or self.obj.room_floor is None:
            return np.array(aabb_low), np.array(aabb_hi)

        # TODO: remove after split floors
        # room_floor will be set to the correct RoomFloor beforehand
        room_instance = self.obj.room_floor.room_instance

        # Get the x-y values from the room segmentation map
        room_aabb_low, room_aabb_hi = self.obj.room_floor.scene.get_aabb_by_room_instance(room_instance)

        if room_aabb_low is None:
            return np.array(aabb_low), np.array(aabb_hi)

        # Use the z values from pybullet
        room_aabb_low[2] = aabb_low[2]
        room_aabb_hi[2] = aabb_hi[2]

        return np.array(room_aabb_low), np.array(room_aabb_hi)

    def _set_value(self, new_value):
        raise NotImplementedError("AABB state currently does not support setting.")

    # Nothing needs to be done to save/load AABB since it will happen due to pose caching.
    def _dump(self):
        return None

    def load(self, data):
        return

[((10, 20, 10, 42), 'igibson.external.pybullet_tools.utils.get_all_links', 'get_all_links', ({(10, 34, 10, 41): 'body_id'}, {}), '(body_id)', False, 'from igibson.external.pybullet_tools.utils import aabb_union, get_aabb, get_all_links\n'), ((12, 28, 12, 45), 'igibson.external.pybullet_tools.utils.aabb_union', 'aabb_union', ({(12, 39, 12, 44): 'aabbs'}, {}), '(aabbs)', False, 'from igibson.external.pybullet_tools.utils import aabb_union, get_aabb, get_all_links\n'), ((11, 17, 11, 45), 'igibson.external.pybullet_tools.utils.get_aabb', 'get_aabb', (), '', False, 'from igibson.external.pybullet_tools.utils import aabb_union, get_aabb, get_all_links\n'), ((31, 15, 31, 38), 'numpy.array', 'np.array', ({(31, 24, 31, 37): 'room_aabb_low'}, {}), '(room_aabb_low)', True, 'import numpy as np\n'), ((31, 40, 31, 62), 'numpy.array', 'np.array', ({(31, 49, 31, 61): 'room_aabb_hi'}, {}), '(room_aabb_hi)', True, 'import numpy as np\n'), ((15, 19, 15, 37), 'numpy.array', 'np.array', ({(15, 28, 15, 36): 'aabb_low'}, {}), '(aabb_low)', True, 'import numpy as np\n'), ((15, 39, 15, 56), 'numpy.array', 'np.array', ({(15, 48, 15, 55): 'aabb_hi'}, {}), '(aabb_hi)', True, 'import numpy as np\n'), ((25, 19, 25, 37), 'numpy.array', 'np.array', ({(25, 28, 25, 36): 'aabb_low'}, {}), '(aabb_low)', True, 'import numpy as np\n'), ((25, 39, 25, 56), 'numpy.array', 'np.array', ({(25, 48, 25, 55): 'aabb_hi'}, {}), '(aabb_hi)', True, 'import numpy as np\n')]