# Path: robel/components/robot/robot.py
# class RobotState:
# class RobotComponent(BaseComponent):
#     def __init__(self,
#                  qpos: Optional[np.ndarray] = None,
#                  qvel: Optional[np.ndarray] = None,
#                  qacc: Optional[np.ndarray] = None):
#     def time(self) -> float:
#     def _process_group(self, **config_kwargs) -> RobotGroupConfig:
#     def step(self,
#              control_groups: Dict[str, np.ndarray],
#              denormalize: bool = True):
#     def set_state(self, state_groups: Dict[str, RobotState], **kwargs):
#     def get_initial_state(
#             self,
#             groups: Union[str, Sequence[str]],
#     ) -> Union[RobotState, Sequence[RobotState]]:
#     def _get_group_states(
#             self, configs: Sequence[RobotGroupConfig]) -> Sequence[RobotState]:
#     def _set_group_states(
#             self, group_states: Sequence[Tuple[RobotGroupConfig, RobotState]]):
#     def _perform_timestep(
#             self,
#             group_controls: Sequence[Tuple[RobotGroupConfig, np.ndarray]]):
#     def _apply_observation_noise(self, state: RobotState,
#                                  config: RobotGroupConfig):
#         def noise(value_range: np.ndarray):
#     def _denormalize_action(self, action: np.ndarray,
#                             config: RobotGroupConfig) -> np.ndarray:
#     def _apply_action_bounds(self, action: np.ndarray,
#                              config: RobotGroupConfig) -> np.ndarray:
# 
# Path: robel/utils/testing/mock_sim_scene.py
# class MockSimScene:
#     """Mock object that implements the SimScene interface."""
# 
#     @staticmethod
#     def create(*unused_args, **unused_kwargs):
#         raise NotImplementedError('`patch_sim_scene` must be called.')
# 
#     def __init__(self, *args, step_duration: float = 1, **kwargs):
#         """Initializes a new mock SimScene."""
#         self.sim = mock.Mock()
#         self.sim.model = MockMjModel(*args, **kwargs)
# 
#         self.model = self.sim.model
#         self.sim.data = self.model
#         self.data = self.model.data
# 
#         self.step_duration = step_duration
# 
#         self.close = mock.Mock()
#         self.advance = mock.Mock()
# 
#         self.renderer = mock.Mock()
#         self.renderer.render_offscreen = lambda w, h, **_: np.zeros((w, h))


from typing import Any
from absl.testing import absltest
from robel.components.robot.robot import ControlMode, RobotComponent
from robel.utils.testing.mock_sim_scene import MockSimScene
import numpy as np



class RobotComponentTest(absltest.TestCase):
    """Unit test class for RobotComponent."""

    def test_get_state(self):
        """Tests querying the state of multiple groups."""
        sim_scene = MockSimScene(nq=10)  # type: Any
        sim_scene.data.qpos[:] = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        sim_scene.data.qvel[:] = [11, 12, 13, 14, 15, 16, 17, 18, 19, 20]
        robot = RobotComponent(
            sim_scene,
            groups={
                'a': {
                    'qpos_indices': [0, 1, 3, 5],
                },
                'b': {
                    'qpos_indices': [2, 6],
                    'qvel_indices': [7, 8, 9],
                },
            })
        a_state, b_state = robot.get_state(['a', 'b'])
        np.testing.assert_array_equal(a_state.qpos, [1, 2, 4, 6])
        np.testing.assert_array_equal(a_state.qvel, [11, 12, 14, 16])
        np.testing.assert_array_equal(b_state.qpos, [3, 7])
        np.testing.assert_array_equal(b_state.qvel, [18, 19, 20])

    def test_step(self):
        """Tests stepping with an action for multiple groups."""

# Path: robel/components/robot/robot.py
# class RobotState:
# class RobotComponent(BaseComponent):
#     def __init__(self,
#                  qpos: Optional[np.ndarray] = None,
#                  qvel: Optional[np.ndarray] = None,
#                  qacc: Optional[np.ndarray] = None):
#     def time(self) -> float:
#     def _process_group(self, **config_kwargs) -> RobotGroupConfig:
#     def step(self,
#              control_groups: Dict[str, np.ndarray],
#              denormalize: bool = True):
#     def set_state(self, state_groups: Dict[str, RobotState], **kwargs):
#     def get_initial_state(
#             self,
#             groups: Union[str, Sequence[str]],
#     ) -> Union[RobotState, Sequence[RobotState]]:
#     def _get_group_states(
#             self, configs: Sequence[RobotGroupConfig]) -> Sequence[RobotState]:
#     def _set_group_states(
#             self, group_states: Sequence[Tuple[RobotGroupConfig, RobotState]]):
#     def _perform_timestep(
#             self,
#             group_controls: Sequence[Tuple[RobotGroupConfig, np.ndarray]]):
#     def _apply_observation_noise(self, state: RobotState,
#                                  config: RobotGroupConfig):
#         def noise(value_range: np.ndarray):
#     def _denormalize_action(self, action: np.ndarray,
#                             config: RobotGroupConfig) -> np.ndarray:
#     def _apply_action_bounds(self, action: np.ndarray,
#                              config: RobotGroupConfig) -> np.ndarray:
# 
# Path: robel/utils/testing/mock_sim_scene.py
# class MockSimScene:
#     """Mock object that implements the SimScene interface."""
# 
#     @staticmethod
#     def create(*unused_args, **unused_kwargs):
#         raise NotImplementedError('`patch_sim_scene` must be called.')
# 
#     def __init__(self, *args, step_duration: float = 1, **kwargs):
#         """Initializes a new mock SimScene."""
#         self.sim = mock.Mock()
#         self.sim.model = MockMjModel(*args, **kwargs)
# 
#         self.model = self.sim.model
#         self.sim.data = self.model
#         self.data = self.model.data
# 
#         self.step_duration = step_duration
# 
#         self.close = mock.Mock()
#         self.advance = mock.Mock()
# 
#         self.renderer = mock.Mock()
#         self.renderer.render_offscreen = lambda w, h, **_: np.zeros((w, h))


# Path: robel/components/robot/robot_test.py
from typing import Any
from absl.testing import absltest
from robel.components.robot.robot import ControlMode, RobotComponent
from robel.utils.testing.mock_sim_scene import MockSimScene
import numpy as np




class RobotComponentTest(absltest.TestCase):
    """Unit test class for RobotComponent."""

    def test_get_state(self):
        """Tests querying the state of multiple groups."""
        sim_scene = MockSimScene(nq=10)  # type: Any
        sim_scene.data.qpos[:] = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        sim_scene.data.qvel[:] = [11, 12, 13, 14, 15, 16, 17, 18, 19, 20]
        robot = RobotComponent(
            sim_scene,
            groups={
                'a': {
                    'qpos_indices': [0, 1, 3, 5],
                },
                'b': {
                    'qpos_indices': [2, 6],
                    'qvel_indices': [7, 8, 9],
                },
            })
        a_state, b_state = robot.get_state(['a', 'b'])
        np.testing.assert_array_equal(a_state.qpos, [1, 2, 4, 6])
        np.testing.assert_array_equal(a_state.qvel, [11, 12, 14, 16])
        np.testing.assert_array_equal(b_state.qpos, [3, 7])
        np.testing.assert_array_equal(b_state.qvel, [18, 19, 20])

    def test_step(self):
        """Tests stepping with an action for multiple groups."""

        sim_scene = MockSimScene(nq=10, ctrl_range=[-1, 1])  # type: Any