"""Tests for the G1 action space."""

import numpy as np
import pytest

from neon.data.action_space import ControlMode, G1ActionSpace, G1_JOINTS


class TestG1ActionSpace:
    """Test the G1 action space definition."""

    def test_joints_count(self):
        """G1 should have 29 joints."""
        assert len(G1_JOINTS) == 29

    def test_joint_indices_unique(self):
        """All joint indices should be unique."""
        indices = [j.index for j in G1_JOINTS]
        assert len(set(indices)) == 29
        assert sorted(indices) == list(range(29))

    def test_arms_only_mode(self):
        """Arms only mode should have 14 joints."""
        space = G1ActionSpace(mode=ControlMode.ARMS_ONLY, include_locomotion=False)
        assert space.num_joints == 14
        assert space.action_dim == 14

    def test_arms_with_locomotion(self):
        """Arms + locomotion should have 14 + 3 = 17 dims."""
        space = G1ActionSpace(mode=ControlMode.ARMS_ONLY, include_locomotion=True)
        assert space.num_joints == 14
        assert space.action_dim == 17

    def test_upper_body_mode(self):
        """Upper body mode: 14 arms + 1 torso + 2 head = 17 joints."""
        space = G1ActionSpace(mode=ControlMode.UPPER_BODY, include_locomotion=False)
        assert space.num_joints == 17

    def test_whole_body_mode(self):
        """Whole body mode: all 29 joints."""
        space = G1ActionSpace(mode=ControlMode.WHOLE_BODY, include_locomotion=False)
        assert space.num_joints == 29
        assert space.action_dim == 29

    def test_whole_body_with_locomotion(self):
        """Whole body + locomotion = 29 + 3 = 32 dims."""
        space = G1ActionSpace(mode=ControlMode.WHOLE_BODY, include_locomotion=True)
        assert space.action_dim == 32

    def test_normalize_denormalize_roundtrip(self):
        """Normalize then denormalize should be identity."""
        space = G1ActionSpace(mode=ControlMode.ARMS_ONLY, include_locomotion=True)
        original = space.default_position()

        normalized = space.normalize(original)
        recovered = space.denormalize(normalized)

        np.testing.assert_allclose(recovered, original, atol=1e-5)

    def test_normalize_range(self):
        """Normalized actions should be in [-1, 1]."""
        space = G1ActionSpace(mode=ControlMode.WHOLE_BODY, include_locomotion=True)
        # Random actions within limits
        rng = np.random.RandomState(42)
        actions = rng.uniform(-1, 1, size=space.action_dim).astype(np.float32)

        normalized = space.normalize(actions)
        assert np.all(normalized >= -1.0)
        assert np.all(normalized <= 1.0)

    def test_split_action(self):
        """Split should produce correct groups."""
        space = G1ActionSpace(mode=ControlMode.ARMS_ONLY, include_locomotion=True)
        action = np.zeros(space.action_dim, dtype=np.float32)
        action[0] = 1.0  # Left arm first joint
        action[7] = 2.0  # Right arm first joint

        split = space.split_action(action)
        assert "left_arm" in split
        assert "right_arm" in split
        assert "locomotion" in split
        assert split["left_arm"][0] == 1.0
        assert split["right_arm"][0] == 2.0

    def test_default_position(self):
        """Default position should have correct shape."""
        space = G1ActionSpace(mode=ControlMode.ARMS_ONLY, include_locomotion=True)
        default = space.default_position()
        assert default.shape == (space.action_dim,)

    def test_serialization(self):
        """to_dict/from_dict should roundtrip."""
        space = G1ActionSpace(mode=ControlMode.UPPER_BODY, include_locomotion=True)
        d = space.to_dict()
        recovered = G1ActionSpace.from_dict(d)
        assert recovered.mode == space.mode
        assert recovered.include_locomotion == space.include_locomotion
        assert recovered.action_dim == space.action_dim

    def test_group_indices(self):
        """Group indices should be valid."""
        space = G1ActionSpace(mode=ControlMode.WHOLE_BODY, include_locomotion=False)
        left_arm = space.get_group_indices("left_arm")
        right_arm = space.get_group_indices("right_arm")
        assert len(left_arm) == 7
        assert len(right_arm) == 7
        # No overlap
        assert set(left_arm).isdisjoint(set(right_arm))
