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, 'igibson.external.pybullet_tools.utils.get_all_links', 'get_all_links', (['body_id'], {}), False, 'from igibson.external.pybullet_tools.utils import aabb_union, get_aabb, get_all_links\n'), (12, 'igibson.external.pybullet_tools.utils.aabb_union', 'aabb_union', (['aabbs'], {}), False, 'from igibson.external.pybullet_tools.utils import aabb_union, get_aabb, get_all_links\n'), (11, 'igibson.external.pybullet_tools.utils.get_aabb', 'get_aabb', (['body_id'], {'link': 'link'}), False, 'from igibson.external.pybullet_tools.utils import aabb_union, get_aabb, get_all_links\n'), (31, 'numpy.array', 'np.array', (['room_aabb_low'], {}), True, 'import numpy as np\n'), (31, 'numpy.array', 'np.array', (['room_aabb_hi'], {}), True, 'import numpy as np\n'), (15, 'numpy.array', 'np.array', (['aabb_low'], {}), True, 'import numpy as np\n'), (15, 'numpy.array', 'np.array', (['aabb_hi'], {}), True, 'import numpy as np\n'), (25, 'numpy.array', 'np.array', (['aabb_low'], {}), True, 'import numpy as np\n'), (25, 'numpy.array', 'np.array', (['aabb_hi'], {}), True, 'import numpy as np\n')]