import numpy as np

from environments.kitchen.spirl.data.block_stacking.src.robosuite.models.base import MujocoXML
from environments.kitchen.spirl.data.block_stacking.src.robosuite.utils.mjcf_utils import (
    array_to_string,
    string_to_array,
)
from environments.kitchen.spirl.data.block_stacking.src.robosuite.utils.mjcf_utils import (
    new_geom,
    new_body,
    new_joint,
)


class Arena(MujocoXML):
    """Base arena class."""

    def set_origin(self, offset):
        """Applies a constant offset to all objects."""
        offset = np.array(offset)
        for node in self.worldbody.findall("./*[@pos]"):
            cur_pos = string_to_array(node.get("pos"))
            new_pos = cur_pos + offset
            node.set("pos", array_to_string(new_pos))

    def add_pos_indicator(self):
        """Adds a new position indicator."""
        body = new_body(name="pos_indicator")
        body.append(
            new_geom(
                "sphere",
                [0.03],
                rgba=[1, 0, 0, 0.5],
                group=1,
                contype="0",
                conaffinity="0",
            )
        )
        body.append(new_joint(type="free", name="pos_indicator"))
        self.worldbody.append(body)
