"""
G1 Action Space — 29 DoF humanoid action definition.

The Unitree G1 humanoid has 29 controllable joints organized into groups.
This module defines the action space, joint limits, and normalization.

Reference: Unitree G1 SDK + GR00T-WholeBodyControl embodiment configs.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, List, Optional, Tuple

import numpy as np


class ControlMode(Enum):
    """Control modes for the G1."""

    ARMS_ONLY = "arms_only"  # 14 DoF arms + 2 grippers + 3 locomotion vel
    UPPER_BODY = "upper_body"  # Arms + torso + head (19 DoF)
    WHOLE_BODY = "whole_body"  # All 29 DoF


@dataclass
class JointInfo:
    """Information about a single joint."""

    name: str
    index: int  # Index in the full 29-DoF vector
    group: str  # "left_arm", "right_arm", "torso", "head", "left_leg", "right_leg"
    lower_limit: float  # Radians
    upper_limit: float  # Radians
    default_pos: float = 0.0  # Default standing position


# fmt: off
# G1 joint definitions (29 joints)
# Limits from Unitree G1 URDF + GR00T-WholeBodyControl configs
G1_JOINTS: List[JointInfo] = [
    # Left Arm (7 DoF) — indices 0-6
    JointInfo("left_shoulder_pitch",  0, "left_arm",  -2.87, 2.87, 0.0),
    JointInfo("left_shoulder_roll",   1, "left_arm",  -0.34, 3.11, 0.3),
    JointInfo("left_shoulder_yaw",    2, "left_arm",  -1.30, 2.18, 0.0),
    JointInfo("left_elbow",           3, "left_arm",  -1.25, 2.61, -0.75),
    JointInfo("left_wrist_yaw",       4, "left_arm",  -1.57, 1.57, 0.0),
    JointInfo("left_wrist_roll",      5, "left_arm",  -0.52, 0.52, 0.0),
    JointInfo("left_gripper",         6, "left_arm",   0.00, 1.00, 0.0),  # 0=open, 1=closed

    # Right Arm (7 DoF) — indices 7-13
    JointInfo("right_shoulder_pitch", 7, "right_arm", -2.87, 2.87, 0.0),
    JointInfo("right_shoulder_roll",  8, "right_arm", -3.11, 0.34, -0.3),
    JointInfo("right_shoulder_yaw",   9, "right_arm", -2.18, 1.30, 0.0),
    JointInfo("right_elbow",          10, "right_arm", -2.61, 1.25, 0.75),
    JointInfo("right_wrist_yaw",      11, "right_arm", -1.57, 1.57, 0.0),
    JointInfo("right_wrist_roll",     12, "right_arm", -0.52, 0.52, 0.0),
    JointInfo("right_gripper",        13, "right_arm",  0.00, 1.00, 0.0),

    # Torso (1 DoF) — index 14
    JointInfo("waist_yaw",            14, "torso",    -2.35, 2.35, 0.0),

    # Head (2 DoF) — indices 15-16
    JointInfo("head_pitch",           15, "head",     -0.52, 0.52, 0.0),
    JointInfo("head_yaw",             16, "head",     -1.57, 1.57, 0.0),

    # Left Leg (6 DoF) — indices 17-22
    JointInfo("left_hip_roll",        17, "left_leg", -0.52, 0.52, 0.0),
    JointInfo("left_hip_yaw",         18, "left_leg", -0.52, 0.52, 0.0),
    JointInfo("left_hip_pitch",       19, "left_leg", -1.57, 0.52, -0.3),
    JointInfo("left_knee",            20, "left_leg",  0.00, 2.53, 0.6),
    JointInfo("left_ankle_pitch",     21, "left_leg", -0.87, 0.52, -0.3),
    JointInfo("left_ankle_roll",      22, "left_leg", -0.17, 0.17, 0.0),

    # Right Leg (6 DoF) — indices 23-28
    JointInfo("right_hip_roll",       23, "right_leg", -0.52, 0.52, 0.0),
    JointInfo("right_hip_yaw",        24, "right_leg", -0.52, 0.52, 0.0),
    JointInfo("right_hip_pitch",      25, "right_leg", -1.57, 0.52, -0.3),
    JointInfo("right_knee",           26, "right_leg",  0.00, 2.53, 0.6),
    JointInfo("right_ankle_pitch",    27, "right_leg", -0.87, 0.52, -0.3),
    JointInfo("right_ankle_roll",     28, "right_leg", -0.17, 0.17, 0.0),
]
# fmt: on


@dataclass
class G1ActionSpace:
    """G1 action space with normalization and group-based access.

    Provides:
    - Joint-level access by name or index
    - Group-level access (left_arm, right_arm, etc.)
    - Normalization to [-1, 1] based on joint limits
    - Mode-based masking (arms_only, upper_body, whole_body)
    - Locomotion velocity (vx, vy, omega) as separate output
    """

    mode: ControlMode = ControlMode.ARMS_ONLY
    include_locomotion: bool = True  # 3 extra dims for (vx, vy, omega)

    # Locomotion limits (m/s, m/s, rad/s)
    locomotion_limits: Tuple[float, float, float] = (1.5, 0.5, 1.0)

    def __post_init__(self):
        self._joints = {j.name: j for j in G1_JOINTS}
        self._groups: Dict[str, List[JointInfo]] = {}
        for j in G1_JOINTS:
            self._groups.setdefault(j.group, []).append(j)

    @property
    def num_joints(self) -> int:
        """Number of controlled joints for current mode."""
        return len(self.active_joints)

    @property
    def action_dim(self) -> int:
        """Total action dimension (joints + optional locomotion)."""
        dim = self.num_joints
        if self.include_locomotion:
            dim += 3  # vx, vy, omega
        return dim

    @property
    def active_joints(self) -> List[JointInfo]:
        """Joints active in the current control mode."""
        if self.mode == ControlMode.ARMS_ONLY:
            return self._groups.get("left_arm", []) + self._groups.get("right_arm", [])
        elif self.mode == ControlMode.UPPER_BODY:
            return (
                self._groups.get("left_arm", [])
                + self._groups.get("right_arm", [])
                + self._groups.get("torso", [])
                + self._groups.get("head", [])
            )
        else:  # WHOLE_BODY
            return G1_JOINTS

    @property
    def active_indices(self) -> List[int]:
        """Indices of active joints in the full 29-DoF vector."""
        return [j.index for j in self.active_joints]

    @property
    def joint_names(self) -> List[str]:
        """Names of active joints."""
        return [j.name for j in self.active_joints]

    def normalize(self, actions: np.ndarray) -> np.ndarray:
        """Normalize raw joint positions to [-1, 1].

        Args:
            actions: Raw joint positions, shape (..., action_dim)

        Returns:
            Normalized actions in [-1, 1], same shape
        """
        result = np.zeros_like(actions)
        joints = self.active_joints
        n_joints = len(joints)

        for i, joint in enumerate(joints):
            lo, hi = joint.lower_limit, joint.upper_limit
            mid = (lo + hi) / 2
            half_range = (hi - lo) / 2
            if half_range > 0:
                result[..., i] = (actions[..., i] - mid) / half_range
            else:
                result[..., i] = 0.0

        # Normalize locomotion if present
        if self.include_locomotion and actions.shape[-1] > n_joints:
            for k, limit in enumerate(self.locomotion_limits):
                idx = n_joints + k
                if limit > 0:
                    result[..., idx] = actions[..., idx] / limit

        return np.clip(result, -1.0, 1.0)

    def denormalize(self, actions: np.ndarray) -> np.ndarray:
        """Denormalize [-1, 1] actions to raw joint positions.

        Args:
            actions: Normalized actions in [-1, 1], shape (..., action_dim)

        Returns:
            Raw joint positions, same shape
        """
        result = np.zeros_like(actions)
        joints = self.active_joints
        n_joints = len(joints)

        for i, joint in enumerate(joints):
            lo, hi = joint.lower_limit, joint.upper_limit
            mid = (lo + hi) / 2
            half_range = (hi - lo) / 2
            result[..., i] = actions[..., i] * half_range + mid

        # Denormalize locomotion
        if self.include_locomotion and actions.shape[-1] > n_joints:
            for k, limit in enumerate(self.locomotion_limits):
                idx = n_joints + k
                result[..., idx] = actions[..., idx] * limit

        return result

    def default_position(self) -> np.ndarray:
        """Get default standing position for active joints."""
        joints = self.active_joints
        pos = np.array([j.default_pos for j in joints], dtype=np.float32)
        if self.include_locomotion:
            pos = np.concatenate([pos, np.zeros(3, dtype=np.float32)])
        return pos

    def get_group_indices(self, group: str) -> List[int]:
        """Get indices within the active action vector for a joint group.

        Args:
            group: One of "left_arm", "right_arm", "torso", "head", "left_leg", "right_leg"

        Returns:
            List of indices within the active action vector
        """
        active_names = self.joint_names
        group_joints = self._groups.get(group, [])
        return [active_names.index(j.name) for j in group_joints if j.name in active_names]

    def split_action(self, action: np.ndarray) -> Dict[str, np.ndarray]:
        """Split a full action vector into named groups.

        Returns dict with keys: left_arm, right_arm, torso, head, left_leg, right_leg, locomotion
        """
        result = {}
        for group_name in ["left_arm", "right_arm", "torso", "head", "left_leg", "right_leg"]:
            indices = self.get_group_indices(group_name)
            if indices:
                result[group_name] = action[..., indices]

        if self.include_locomotion:
            n = self.num_joints
            result["locomotion"] = action[..., n : n + 3]

        return result

    def to_dict(self) -> Dict:
        """Serialize to dict for config files."""
        return {
            "mode": self.mode.value,
            "include_locomotion": self.include_locomotion,
            "locomotion_limits": list(self.locomotion_limits),
            "num_joints": self.num_joints,
            "action_dim": self.action_dim,
            "joint_names": self.joint_names,
        }

    @classmethod
    def from_dict(cls, d: Dict) -> "G1ActionSpace":
        """Deserialize from dict."""
        return cls(
            mode=ControlMode(d["mode"]),
            include_locomotion=d.get("include_locomotion", True),
            locomotion_limits=tuple(d.get("locomotion_limits", (1.5, 0.5, 1.0))),
        )

    def __repr__(self) -> str:
        return (
            f"G1ActionSpace(mode={self.mode.value}, "
            f"joints={self.num_joints}, action_dim={self.action_dim})"
        )
