# ----------------------------------------------------------------------------
# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ----------------------------------------------------------------------------

import torch
import numpy as np

from typing import List, Dict, Tuple, Union, Sequence
from dataclasses import dataclass, field
from tensordict import TensorDict

from dexsim.engine import Articulation as _Articulation
from embodichain.lab.sim.cfg import RobotCfg
from embodichain.lab.sim.solvers import SolverCfg, BaseSolver
from embodichain.lab.sim.objects import Articulation
from embodichain.lab.sim.utility.tensor import to_tensor
from embodichain.utils.math import quat_from_matrix, matrix_from_quat
from embodichain.utils.string import (
    is_regular_expression,
    resolve_matching_names_values,
)
from embodichain.utils import logger


@dataclass
class ControlGroup:
    r"""Represents a group of controllable joints in a robot.

    Attributes:
        joint_names (List[str]): Names of the joints in this control group.
        joint_ids (List[int]): IDs corresponding to the joints in this control group.
        link_names (List[str]): Names of child links associated with the joints.
    """

    joint_names: List[str] = field(default_factory=list)
    joint_ids: List[int] = field(default_factory=list)
    link_names: List[str] = field(default_factory=list)

    def __post_init__(self):
        pass


class Robot(Articulation):
    """A class representing a batch of robots in the simulation environment.

    Robot is a specific type of articulation that can have additional properties or methods.
    - `control_parts`: Specify the parts that can be controlled in a different manner. Different part may have
        different joint ids, drive properties, pyhsical attributes, kinematic solvers or motion planners.
    - `solvers`: Specify the kinematic solvers for the robot.
    - `planners`: Specify the motion planner for the robot.
    """

    def __init__(
        self,
        cfg: RobotCfg,
        entities: List[_Articulation],
        device: torch.device = torch.device("cpu"),
    ) -> None:

        self._entities = entities
        self.cfg = cfg

        # Initialize joint ids for control parts.
        self._joint_ids: Dict[str, List[int]] = {}

        self._control_groups: Dict[str, ControlGroup] = {}

        if self.cfg.control_parts:
            self._init_control_parts(self.cfg.control_parts)

        super().__init__(cfg, entities, device)

        self._solvers = {}

        if self.cfg.solver_cfg:
            self.init_solver(self.cfg.solver_cfg)

    def __str__(self) -> str:
        parent_str = super().__str__()
        return (
            parent_str
            + f" | control_parts: {self.control_parts}, solvers: {self._solvers}"
        )

    @property
    def control_parts(self) -> Dict[str, List[str]] | None:
        """Get the control parts of the robot."""
        return self.cfg.control_parts

    def get_joint_ids(
        self, name: str | None = None, remove_mimic: bool = False
    ) -> List[int]:
        """Get the joint ids of the robot for a specific control part.

        Args:
            name (str | None): The name of the control part to get the joint ids for. If None, the default part is used.
            remove_mimic (bool, optional): If True, mimic joints will be excluded from the returned joint ids. Defaults to False.

        Returns:
            List[int]: The joint ids of the robot for the specified control part.
        """
        if not self.control_parts or name is None:
            return (
                torch.arange(self.dof, dtype=torch.int32).tolist()
                if not remove_mimic
                else self.active_joint_ids
            )

        if name not in self.control_parts:
            logger.log_error(
                f"The control part '{name}' does not exist in the robot's control parts."
            )
        return (
            self._joint_ids[name]
            if not remove_mimic
            else [i for i in self._joint_ids[name] if i not in self.mimic_ids]
        )

    def get_link_names(self, name: str | None = None) -> Union[List[str], None]:
        """Get the link names of the robot for a specific control part.

        If no control part is specified, return all link names.

        Args:
            name (str, optional): The name of the control part to get the link names for. If None, the default part is used.

        Returns:
            List[str]: The link names of the robot for the specified control part.
        """
        if not self.control_parts or name is None:
            return self.link_names

        if name not in self.control_parts:
            logger.log_error(
                f"The control part '{name}' does not exist in the robot's control parts {self.control_parts}."
            )
        return self._control_groups[name].link_names

    def get_qpos_limits(
        self, name: str | None = None, env_ids: Sequence[int] | None = None
    ) -> torch.Tensor:
        """Get the joint position limits (qpos) of the robot for a specific control part.

        It returns all joint position limits if no control part is specified.

        Args:
            name (str | None): The name of the control part to get the qpos limits for.
            env_ids (Sequence[int] | None): The environment ids to get the qpos limits for. If None, all environments are used.

        Returns:
            torch.Tensor: Joint position limits with shape (N, dof, 2), where N is the number of environments.
        """
        local_env_ids = self._all_indices if env_ids is None else env_ids

        qpos_limits = self.body_data.qpos_limits
        if name is None:
            return qpos_limits[local_env_ids, :]
        else:
            if not self.control_parts or name not in self.control_parts:
                logger.log_error(
                    f"The control part '{name}' does not exist in the robot's control parts."
                )
            part_joint_ids = self.get_joint_ids(name=name)
            return qpos_limits[local_env_ids][:, part_joint_ids, :]

    def get_qvel_limits(
        self, name: str | None = None, env_ids: Sequence[int] | None = None
    ) -> torch.Tensor:
        """Get the joint velocity limits (qvel) of the robot for a specific control part.

        It returns all joint velocity limits if no control part is specified.

        Args:
            name (str | None): The name of the control part to get the qvel limits for.
            env_ids (Sequence[int] | None): The environment ids to get the qvel limits for. If None, all environments are used.

        Returns:
            torch.Tensor: Joint velocity limits with shape (N, dof), where N is the number of environments.
        """
        local_env_ids = self._all_indices if env_ids is None else env_ids

        qvel_limits = self.body_data.qvel_limits
        if name is None:
            return qvel_limits[local_env_ids, :]
        else:
            if not self.control_parts or name not in self.control_parts:
                logger.log_error(
                    f"The control part '{name}' does not exist in the robot's control parts."
                )
            part_joint_ids = self.get_joint_ids(name=name)
            return qvel_limits[local_env_ids][:, part_joint_ids]

    def get_qf_limits(
        self, name: str | None = None, env_ids: Sequence[int] | None = None
    ) -> torch.Tensor:
        """Get the joint effort limits (qf) of the robot for a specific control part.

        It returns all joint effort limits if no control part is specified.

        Args:
            name (str | None): The name of the control part to get the qf limits for.
            env_ids (Sequence[int] | None): The environment ids to get the qf limits for. If None, all environments are used.

        Returns:
            torch.Tensor: Joint effort limits with shape (N, dof), where N is the number of environments.
        """
        local_env_ids = self._all_indices if env_ids is None else env_ids

        qf_limits = self.body_data.qf_limits
        if name is None:
            return qf_limits[local_env_ids, :]
        else:
            if not self.control_parts or name not in self.control_parts:
                logger.log_error(
                    f"The control part '{name}' does not exist in the robot's control parts."
                )
            part_joint_ids = self.get_joint_ids(name=name)
            return qf_limits[local_env_ids][:, part_joint_ids]

    def get_proprioception(self) -> TensorDict[str, torch.Tensor]:
        """Gets robot proprioception information, primarily for agent state representation in robot learning scenarios.

        The default proprioception information includes:
            - qpos: Joint positions.
            - qvel: Joint velocities.
            - qf: Joint efforts.

        Returns:
            Dict[str, torch.Tensor]: A dictionary containing the robot's proprioception information
        """

        return TensorDict(
            qpos=self.body_data.qpos,
            qvel=self.body_data.qvel,
            qf=self.body_data.qf,
            batch_size=[self.num_instances],
            device=self.device,
        )

    def set_qpos(
        self,
        qpos: torch.Tensor,
        joint_ids: Sequence[int] | None = None,
        env_ids: Sequence[int] | None = None,
        target: bool = True,
        name: str | None = None,
    ) -> None:
        """Set the joint positions (qpos) or target positions for the articulation.

        Args:
            qpos (torch.Tensor): Joint positions with shape (N, dof), where N is the number of environments.
            joint_ids (Sequence[int] | None, optional): Joint indices to apply the positions. If None, applies to all joints.
            env_ids (Sequence[int] | None): Environment indices to apply the positions. Defaults to all environments.
            target (bool): If True, sets target positions for simulation. If False, updates current positions directly.
            name (str | None): The name of the control part to set the qpos for. If None, the default part is used.

        Raises:
            ValueError: If the length of `env_ids` does not match the length of `qpos`.
        """
        if name is None:
            super().set_qpos(
                qpos=qpos,
                joint_ids=joint_ids,
                env_ids=env_ids,
                target=target,
            )
        else:
            if not self.control_parts or name not in self.control_parts:
                logger.log_error(
                    f"The control part '{name}' does not exist in the robot's control parts."
                )
            part_joint_ids = self.get_joint_ids(name=name)
            if joint_ids is not None:
                logger.log_warning(f"`joint_ids` is ignored when `name` is specified.")

            super().set_qpos(
                qpos=qpos,
                joint_ids=part_joint_ids,
                env_ids=env_ids,
                target=target,
            )

    def get_qpos(self, name: str | None = None, target: bool = False) -> torch.Tensor:
        """Get the joint positions (qpos) of the robot.

        Args:
            name (str | None): The name of the control part to get the qpos for. If None, the default part is used.
            target (bool): If True, gets target positions for simulation. If False, gets current positions. The default is False.

        Returns:
            torch.Tensor: Joint positions with shape (N, dof), where N is the number of environments.
        """

        qpos = super().get_qpos(target=target)
        if name is None:
            return qpos
        else:
            if not self.control_parts or name not in self.control_parts:
                logger.log_error(
                    f"The control part '{name}' does not exist in the robot's control parts."
                )
            part_joint_ids = self.get_joint_ids(name=name)
            return qpos[:, part_joint_ids]

    def set_qvel(
        self,
        qvel: torch.Tensor,
        joint_ids: Sequence[int] | None = None,
        env_ids: Sequence[int] | None = None,
        target: bool = True,
        name: str | None = None,
    ) -> None:
        """Set the joint velocities (qvel) or target velocities for the articulation.

        Args:
            qvel (torch.Tensor): Joint velocities with shape (N, dof), where N is the number of environments.
            joint_ids (Sequence[int] | None, optional): Joint indices to apply the velocities. If None, applies to all joints.
            env_ids (Sequence[int] | None): Environment indices to apply the velocities. Defaults to all environments.
            target (bool): If True, sets target velocities for simulation. If False, updates current velocities directly.
            name (str | None): The name of the control part to set the qvel for. If None, the default part is used.

        Raises:
            ValueError: If the length of `env_ids` does not match the length of `qvel`.
        """
        if name is None:
            super().set_qvel(
                qvel=qvel,
                joint_ids=joint_ids,
                env_ids=env_ids,
                target=target,
            )
        else:
            if not self.control_parts or name not in self.control_parts:
                logger.log_error(
                    f"The control part '{name}' does not exist in the robot's control parts."
                )
            part_joint_ids = self.get_joint_ids(name=name)
            if joint_ids is not None:
                logger.log_warning(f"`joint_ids` is ignored when `name` is specified.")

            super().set_qvel(
                qvel=qvel,
                joint_ids=part_joint_ids,
                env_ids=env_ids,
                target=target,
            )

    def get_qvel(self, name: str | None = None, target: bool = False) -> torch.Tensor:
        """Get the joint velocities (qvel) of the robot.

        Args:
            name (str | None): The name of the control part to get the qvel for. If None, the default part is used.
            target (bool): If True, gets target velocities for simulation. If False, gets current velocities. The default is False.

        Returns:
            torch.Tensor: Joint velocities with shape (N, dof), where N is the number of environments.
        """

        qvel = super().get_qvel(target=target)
        if name is None:
            return qvel
        else:
            if not self.control_parts or name not in self.control_parts:
                logger.log_error(
                    f"The control part '{name}' does not exist in the robot's control parts."
                )
            part_joint_ids = self.get_joint_ids(name=name)
            return qvel[:, part_joint_ids]

    def set_qf(
        self,
        qf: torch.Tensor,
        joint_ids: Sequence[int] | None = None,
        env_ids: Sequence[int] | None = None,
        name: str | None = None,
    ) -> None:
        """Set the joint efforts (qf) for the articulation.

        Args:
            qf (torch.Tensor): Joint efforts with shape (N, dof), where N is the number of environments.
            joint_ids (Sequence[int] | None, optional): Joint indices to apply the efforts. If None, applies to all joints.
            env_ids (Sequence[int] | None): Environment indices to apply the efforts. Defaults to all environments.
            name (str | None): The name of the control part to set the qf for. If None, the default part is used.

        Raises:
            ValueError: If the length of `env_ids` does not match the length of `qf`.
        """
        if name is None:
            super().set_qf(
                qf=qf,
                joint_ids=joint_ids,
                env_ids=env_ids,
            )
        else:
            if not self.control_parts or name not in self.control_parts:
                logger.log_error(
                    f"The control part '{name}' does not exist in the robot's control parts."
                )
            part_joint_ids = self.get_joint_ids(name=name)
            if joint_ids is not None:
                logger.log_warning(f"`joint_ids` is ignored when `name` is specified.")

            super().set_qf(
                qf=qf,
                joint_ids=part_joint_ids,
                env_ids=env_ids,
            )

    def get_qf(self, name: str | None = None) -> torch.Tensor:
        """Get the joint efforts (qf) of the robot.

        Args:
            name (str | None): The name of the control part to get the qf for. If None, the default part is used.
        Returns:
            torch.Tensor: Joint efforts with shape (N, dof), where N is the number of environments.
        """

        qf = super().get_qf()
        if name is None:
            return qf
        else:
            if not self.control_parts or name not in self.control_parts:
                logger.log_error(
                    f"The control part '{name}' does not exist in the robot's control parts."
                )
            part_joint_ids = self.get_joint_ids(name=name)
            return qf[:, part_joint_ids]

    def compute_fk(
        self,
        qpos: torch.Tensor | np.ndarray | None,
        name: str | None = None,
        link_names: List[str] | None = None,
        end_link_name: str | None = None,
        root_link_name: str | None = None,
        env_ids: Sequence[int] | None = None,
        to_matrix: bool = False,
    ) -> torch.Tensor:
        """Compute the forward kinematics of the robot given joint positions and optionally a specific part name.
        The output pose will be in the local arena frame.

        Args:
            qpos (torch.Tensor | np.ndarray | None): Joint positions of the robot, (n_envs, num_joints).
            name (str | None): The name of the control part to compute the FK for. If None, the default part is used.
            link_names (List[str] | None): The names of the links to compute the FK for. If None, all links are used.
            end_link_name (str | None): The name of the end link to compute the FK for. If None, the default end link is used.
            root_link_name (str | None): The name of the root link to compute the FK for. If None, the default root link is used.
            env_ids (Sequence[int] | None): The environment ids to compute the FK for. If None, all environments are used.
            to_matrix (bool): If True, returns the transformation in the form of a 4x4 matrix.

        Returns:
            torch.Tensor: The forward kinematics result with shape (n_envs, 7) or (n_envs, 4, 4) if `to_matrix` is True.
        """
        local_env_ids = self._all_indices if env_ids is None else env_ids

        if name is None and hasattr(super(), "compute_fk"):
            return super().compute_fk(
                qpos=qpos,
                link_names=link_names,
                end_link_name=end_link_name,
                root_link_name=root_link_name,
            )

        if not self._solvers:
            logger.log_error(
                "No solvers are defined for the robot. Please ensure that the robot has solvers configured."
            )

        solver = self._solvers.get(name if name is not None else "default", None)
        if solver is None:
            logger.log_error(
                f"The control part '{name}' does not have an associated solver. Please ensure that a valid control part with an available solver is provided."
            )
            return None

        if qpos.dim() == 1:
            qpos = qpos.unsqueeze(0)

        if qpos.shape[0] != len(local_env_ids):
            logger.log_error(
                f"Joint positions batch size mismatch. Expected {len(local_env_ids)} but got {qpos.shape[0]}."
            )

        if qpos.shape[1] != solver.dof:
            logger.log_error(
                f"Joint positions shape mismatch. Expected {solver.dof} joints, got {qpos.shape[1]}."
            )

        result_matrix = solver.get_fk(qpos=qpos)

        base_pose = self.get_link_pose(
            link_name=solver.root_link_name, env_ids=local_env_ids, to_matrix=True
        )
        result_matrix = torch.bmm(base_pose, result_matrix)

        if to_matrix:
            return result_matrix
        else:
            pos = result_matrix[:, :3, 3]
            quat = quat_from_matrix(result_matrix[:, :3, :3])
            return torch.cat((pos, quat), dim=-1)

    def compute_ik(
        self,
        pose: torch.Tensor | np.ndarray,
        joint_seed: torch.Tensor | np.ndarray | None = None,
        name: str | None = None,
        env_ids: Sequence[int] | None = None,
        return_all_solutions: bool = False,
    ) -> Tuple[torch.Tensor, torch.Tensor] | None:
        """Compute the inverse kinematics of the robot given joint positions and optionally a specific part name.
        The input pose should be in the local arena frame.

        Args:
            pose (torch.Tensor): The end effector pose of the robot, (n_envs, 7) or (n_envs, 4, 4).
            joint_seed (torch.Tensor | None): The joint positions to use as a seed for the IK computation, (n_envs, dof).
                If None, the zero joint positions will be used as the seed.
            name (str | None): The name of the control part to compute the IK for. If None, the default part is used.
            env_ids (Sequence[int] | None): Environment indices to apply the positions. Defaults to all environments.
            return_all_solutions (bool): Whether to return all IK solutions or just the best one. Defaults to False.

        Returns:
            Tuple[torch.Tensor, torch.Tensor] | None: The success Tensor with shape (n_envs, ) and qpos Tensor with shape (n_envs, max_results, dof), or None if solver not found.
        """
        local_env_ids = self._all_indices if env_ids is None else env_ids

        solver = self._solvers.get(name if name is not None else "default", None)
        if solver is None:
            logger.log_error(
                f"The control part '{name}' does not have an associated solver. Please ensure that a valid control part with an available solver is provided."
            )
            return None

        pose = to_tensor(pose, device=self.device)
        if (pose.dim() == 1 and pose.shape[1] == 7) or (
            pose.dim() == 2 and pose.shape[1] == 4
        ):
            pose = pose.unsqueeze(0)

        if pose.shape[0] != len(local_env_ids):
            logger.log_error(
                f"Pose batch size mismatch. Expected {len(local_env_ids)} but got {pose.shape[0]}."
            )

        if joint_seed is not None:
            joint_seed = to_tensor(joint_seed, device=self.device)
            if joint_seed.dim() == 1:
                joint_seed = joint_seed.unsqueeze(0)

            if joint_seed.shape[0] != len(local_env_ids):
                logger.log_error(
                    f"Joint seed batch size mismatch. Expected {len(local_env_ids)} but got {joint_seed.shape[0]}."
                )

        if pose.shape[-1] == 7 and pose.dim() == 2:
            # Convert pose from (batch, 7) to (batch, 4, 4)
            pos = pose[:, :3]
            quat = pose[:, 3:]
            # Convert quaternion to rotation matrix
            rot = matrix_from_quat(quat)
            # Build homogeneous transformation matrix efficiently
            pose = torch.eye(4, device=pose.device).repeat(pose.shape[0], 1, 1)
            pose[:, :3, :3] = rot
            pose[:, :3, 3] = pos

        base_pose = self.get_link_pose(
            link_name=solver.root_link_name, env_ids=local_env_ids, to_matrix=True
        )
        pose = torch.bmm(torch.inverse(base_pose), pose)

        ret, qpos = solver.get_ik(
            target_xpos=pose,
            qpos_seed=joint_seed,
            return_all_solutions=return_all_solutions,
        )
        dof = qpos.shape[-1]
        if not return_all_solutions:
            qpos = qpos.reshape(-1, dof)

        return ret.to(self.device), qpos.to(self.device)

    def compute_batch_fk(
        self,
        qpos: torch.Tensor,
        name: str,
        env_ids: Sequence[int] | None = None,
        to_matrix: bool = False,
    ):
        """Compute the forward kinematics of the robot given joint positions and optionally a specific part name.
        The output pose will be in the local arena frame.

        Args:
            qpos (torch.Tensor | np.ndarray | None): Joint positions of the robot, (n_envs, n_batch, num_joints).
            name (str | None): The name of the control part to compute the FK for. If None, the default part is used.
            env_ids (Sequence[int] | None): The environment ids to compute the FK for. If None, all environments are used.
            to_matrix (bool): If True, returns the transformation in the form of a 4x4 matrix.

        Returns:
            torch.Tensor: The forward kinematics result with shape (n_envs, batch, 7) or (n_envs, batch, 4, 4) if `to_matrix` is True.
        """
        local_env_ids = self._all_indices if env_ids is None else env_ids
        if not self._solvers:
            logger.log_error(
                "No solvers are defined for the robot. Please ensure that the robot has solvers configured."
            )

        solver = self._solvers.get(name if name is not None else "default", None)
        if solver is None:
            logger.log_error(
                f"The control part '{name}' does not have an associated solver. Please ensure that a valid control part with an available solver is provided."
            )
            return None

        if qpos.shape[0] != len(local_env_ids):
            logger.log_error(
                f"Joint positions batch size mismatch. Expected {len(local_env_ids)} but got {qpos.shape[0]}."
            )

        if qpos.shape[2] != solver.dof:
            logger.log_error(
                f"Joint positions shape mismatch. Expected {solver.dof} joints, got {qpos.shape[1]}."
            )

        n_batch = qpos.shape[1]
        qpos_batch = qpos.reshape(-1, solver.dof)
        xpos_batch = solver.get_fk(qpos=qpos_batch)

        # get xpos from link root
        base_xpos_n_envs = self.get_link_pose(
            link_name=solver.root_link_name, env_ids=local_env_ids, to_matrix=True
        )
        base_xpos_batch = (
            base_xpos_n_envs[:, None, :, :].repeat(1, n_batch, 1, 1).reshape(-1, 4, 4)
        )
        result_matrix = torch.bmm(base_xpos_batch, xpos_batch)

        if to_matrix:
            result_matrix = result_matrix.reshape(len(local_env_ids), n_batch, 4, 4)
            return result_matrix
        else:
            pos = result_matrix[:, :3, 3]
            quat = quat_from_matrix(result_matrix[:, :3, :3])
            result = torch.cat((pos, quat), dim=-1)
            result = result.reshape(len(local_env_ids), n_batch, 7)
            return result

    def compute_batch_ik(
        self,
        pose: torch.Tensor | np.ndarray,
        joint_seed: torch.Tensor | np.ndarray | None,
        name: str,
        env_ids: Sequence[int] | None = None,
    ):
        """Compute the inverse kinematics of the robot given joint positions and optionally a specific part name.
        The input pose should be in the local arena frame.

        Args:
            pose (torch.Tensor): The end effector pose of the robot, (n_envs, n_batch, 7) or (n_envs, n_batch, 4, 4).
            joint_seed (torch.Tensor | None): The joint positions to use as a seed for the IK computation, (n_envs, n_batch, dof). If None, the zero joint positions will be used as the seed.
            name (str | None): The name of the control part to compute the IK for. If None, the default part is used.
            env_ids (Sequence[int] | None): Environment indices to apply the positions. Defaults to all environments.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]:
                Success Tensor with shape (n_envs, n_batch)
                Qpos Tensor with shape (n_envs, n_batch, dof).
        """
        local_env_ids = self._all_indices if env_ids is None else env_ids

        solver = self._solvers.get(name if name is not None else "default", None)
        if solver is None:
            logger.log_error(
                f"The control part '{name}' does not have an associated solver. Please ensure that a valid control part with an available solver is provided."
            )
            return None
        pose = to_tensor(pose, device=self.device)

        if pose.shape[0] != len(local_env_ids):
            logger.log_error(
                f"Pose batch size mismatch. Expected {len(local_env_ids)} but got {pose.shape[0]}."
            )

        n_batch = pose.shape[1]
        n_dof = solver.dof
        if joint_seed is None:
            joint_seed = torch.zeros(
                (len(local_env_ids), n_batch, n_dof),
                dtype=torch.float32,
                device=self.device,
            )

        if joint_seed.shape[0] != len(local_env_ids):
            logger.log_error(
                f"Joint seed env size mismatch. Expected {len(local_env_ids)} but got {joint_seed.shape[0]}."
            )

        if joint_seed.shape[1] != n_batch:
            logger.log_error(
                f"Joint seed batch size mismatch. Expected {n_batch} but got {joint_seed.shape[1]}."
            )

        if joint_seed.shape[-1] != n_dof:
            logger.log_error(
                f"Joint seed dof size mismatch. Expected {n_batch} but got {joint_seed.shape[-1]}."
            )

        if pose.shape[-1] == 7 and pose.dim() == 3:
            # Convert pose from (n_envs, n_batch, 7) to (n_envs * n_batch, 4, 4)
            pose_batch = pose.reshape(-1, 7)
            pos = pose_batch[:, :3]
            quat = pose_batch[:, 3:]
            # Convert quaternion to rotation matrix
            rot = matrix_from_quat(quat)
            # Build homogeneous transformation matrix efficiently
            pose_batch = torch.eye(4, device=pose.device).repeat(
                pose_batch.shape[0], 1, 1
            )
            pose_batch[:, :3, :3] = rot
            pose_batch[:, :3, 3] = pos
        else:
            # Convert pose from (n_envs, n_batch, 4, 4) to (n_envs * n_batch, 4, 4)
            pose_batch = pose.reshape(-1, 4, 4)

        # get xpos from link root
        base_xpos_n_envs = self.get_link_pose(
            link_name=solver.root_link_name, env_ids=local_env_ids, to_matrix=True
        )
        base_inv_xpos_n_envs = torch.inverse(base_xpos_n_envs)
        base_inv_xpos_batch = (
            base_inv_xpos_n_envs[:, None, :, :]
            .repeat(1, n_batch, 1, 1)
            .reshape(-1, 4, 4)
        )
        pose_batch = torch.bmm(base_inv_xpos_batch, pose_batch)

        joint_seed_batch = joint_seed.reshape(-1, n_dof)
        ret, qpos_batch = solver.get_ik(
            target_xpos=pose_batch,
            qpos_seed=joint_seed_batch,
            return_all_solutions=False,
        )
        ret = ret.reshape(len(local_env_ids), n_batch)
        qpos = qpos_batch.reshape(len(local_env_ids), n_batch, n_dof)
        return ret, qpos

    def _init_control_parts(self, control_parts: Dict[str, List[str]]) -> None:
        """Initialize the control parts of the robot.

        Args:
            control_parts (Dict[str, List[str]]): A dictionary where keys are control part names and values are lists of
                joint names or regular expressions that match joint names.
        """
        joint_name_to_ids = {
            name: i
            for i, name in enumerate(self._entities[0].get_actived_joint_names())
        }
        for name, joint_names in control_parts.items():
            # convert joint_names which is a regular expression to a list of joint names
            joint_names_expanded = []
            for jn in joint_names:
                if is_regular_expression(jn):
                    _, names, _ = resolve_matching_names_values(
                        {jn: None}, self.joint_names
                    )
                    joint_names_expanded.extend(names)
                else:
                    joint_names_expanded.append(jn)

            self._joint_ids[name] = [
                joint_name_to_ids[joint_name]
                for joint_name in joint_names_expanded
                if joint_name in joint_name_to_ids
            ]
            if len(self._joint_ids[name]) != len(joint_names_expanded):
                logger.log_error(
                    f"joint names in control part '{name}' do not match the robot's joint names. The full joint names are: {self.joint_names}."
                )
            self.cfg.control_parts[name] = joint_names_expanded

        # Initialize control groups
        self._control_groups = self._extract_control_groups()

    def _set_default_joint_drive(self) -> None:
        """Set default joint drive parameters based on the configuration."""
        import numbers
        from embodichain.utils.string import resolve_matching_names_values

        drive_props = [
            ("damping", self.default_joint_damping),
            ("stiffness", self.default_joint_stiffness),
            ("max_effort", self.default_joint_max_effort),
            ("max_velocity", self.default_joint_max_velocity),
            ("friction", self.default_joint_friction),
        ]

        for prop_name, default_array in drive_props:
            value = getattr(self.cfg.drive_pros, prop_name, None)
            if value is None:
                continue
            if isinstance(value, numbers.Number):
                default_array[:] = value
            else:
                try:
                    control_part_dict = {}
                    value_copy = value.copy()
                    if self.control_parts:
                        # Extract control part and map the corresponding joint names
                        for key in value.keys():
                            if key in self.control_parts:
                                control_part_dict[key] = value_copy.pop(key)

                    indices, _, values = resolve_matching_names_values(
                        value_copy, self.joint_names
                    )

                    if self.control_parts:
                        # Add control part joints to indices and values
                        for part_name, part_value in control_part_dict.items():
                            part_joint_names = self.control_parts[part_name]
                            part_indices, _, part_values = (
                                resolve_matching_names_values(
                                    {jn: part_value for jn in part_joint_names},
                                    self.joint_names,
                                )
                            )
                            indices.extend(part_indices)
                            values.extend(part_values)

                    default_array[:, indices] = torch.as_tensor(
                        values, dtype=torch.float32, device=self.device
                    )
                except Exception as e:
                    logger.log_error(f"Failed to set {prop_name}: {e}")

        drive_pros = self.cfg.drive_pros
        if isinstance(drive_pros, dict):
            drive_type = drive_pros.get("drive_type", "force")
        else:
            drive_type = getattr(drive_pros, "drive_type", "force")

        # Apply drive parameters to all articulations in the batch
        self.set_drive(
            stiffness=self.default_joint_stiffness,
            damping=self.default_joint_damping,
            max_effort=self.default_joint_max_effort,
            max_velocity=self.default_joint_max_velocity,
            friction=self.default_joint_friction,
            drive_type=drive_type,
        )

    def init_solver(self, cfg: Union[SolverCfg, Dict[str, SolverCfg]]) -> None:
        """Initialize the kinematic solver for the robot.

        Args:
            cfg (Union[SolverCfg, Dict[str, SolverCfg]]): The configuration for the kinematic solver.
        """
        self.cfg: RobotCfg

        if isinstance(cfg, SolverCfg):
            if self.control_parts:
                logger.log_error(
                    "Control parts are defined in the robot configuration, solver_cfg must be a dictionary."
                )

            if cfg.urdf_path is None:
                cfg.urdf_path = self.cfg.fpath
            self._solvers["default"] = cfg.init_solver(device=self.device)
        elif isinstance(cfg, Dict):
            if isinstance(self.cfg.control_parts, Dict) is False:
                logger.log_error(
                    "When `solver_cfg` is a dictionary, `control_parts` must also be a dictionary."
                )

            # If solver_cfg is a dictionary, iterate through it to create solvers
            for name, solver_cfg in cfg.items():
                if solver_cfg.urdf_path is None:
                    solver_cfg.urdf_path = self.cfg.fpath
                _, part_names, value = resolve_matching_names_values(
                    {name: solver_cfg}, self.cfg.control_parts.keys()
                )
                for part_name in part_names:
                    if (
                        not hasattr(solver_cfg, "joint_names")
                        or solver_cfg.joint_names is None
                    ):
                        solver_cfg.joint_names = self.cfg.control_parts[part_name]
                    self._solvers[name] = solver_cfg.init_solver(device=self.device)

    def get_solver(self, name: str | None = None) -> BaseSolver | None:
        """Get the kinematic solver for a specific control part.

        Args:
            name (str | None): The name of the control part to get the solver for. If None, the default part is used.

        Returns:
            BaseSolver | None: The kinematic solver for the specified control part, or None if not found.
        """

        if not self._solvers:
            logger.log_error(
                "No solvers are defined for the robot. Please ensure that the robot has solvers configured."
            )
            return None

        return self._solvers.get(name if name is not None else "default", None)

    def get_control_part_base_pose(
        self,
        name: str | None = None,
        env_ids: Sequence[int] | None = None,
        to_matrix: bool = False,
    ) -> torch.Tensor:
        """Retrieves the base pose of the control part for a specified robot.

        Args:
            name (str | None): The name of the control part the solver adhere to. If None, the default solver is used.
            env_ids (Sequence[int] | None): A sequence of environment IDs to specify the environments. If None, all indices are used.
            to_matrix (bool): If True, returns the pose in the form of a 4x4 matrix.

        Returns:
            The pose of the specified link in the form of a matrix.
        """
        local_env_ids = self._all_indices if env_ids is None else env_ids

        root_link_name = None
        if name in self._control_groups:
            root_link_name = self._control_groups[name].link_names[0]

        return self.get_link_pose(
            link_name=root_link_name, env_ids=local_env_ids, to_matrix=to_matrix
        )

    def get_control_part_link_names(self, name: str | None = None) -> List[str]:
        """Get the link names of the control part.

        Args:
            name (str | None): The name of the control part. If None, return all link names.
        Returns:
            List[str]: link names of the control part.
        """
        if name is None:
            return self.link_names
        if name in self._control_groups:
            return self._control_groups[name].link_names
        else:
            logger.log_warning(
                f"The control part '{name}' does not exist in the robot's control parts."
            )
            return []

    def _extract_control_groups(self) -> Dict[str, ControlGroup]:
        r"""Extract control groups from the active joint names.

        This method creates a dictionary of control groups where each control
        group is associated with its corresponding joint names. It utilizes
        the `_extract_control_group` method to populate the control groups.

        Returns:
            Dict[str, ControlGroup]: A dictionary mapping control group names
                                    to their corresponding ControlGroup instances.
        """
        if not self.control_parts:
            return {}

        control_groups = {
            control_group_name: self._extract_control_group(joint_names)
            for control_group_name, joint_names in self.control_parts.items()
        }

        return control_groups

    def _extract_control_group(self, joint_names: List[str]) -> ControlGroup:
        r"""Extract a control group from the given list of joint names.

        Args:
            joint_names (List[str]): A list of joint names
                                        to be included in the control group.

        Returns:
            ControlGroup: An instance of ControlGroup containing the specified joints
                            and their associated links.
        """
        control_group = ControlGroup()
        joint_id_list = []

        for joint_name in joint_names:
            if joint_name in self.joint_names:
                joint_index = self.joint_names.index(joint_name)
                joint_id_list.append(joint_index)
                control_group.joint_names.append(joint_name)

                # Set root link for first joint
                if len(control_group.link_names) == 0:
                    parent_names = self._entities[0].get_ancestral_link_names(
                        joint_index
                    )
                    control_group.link_names.extend(parent_names)

                child_name = self._entities[0].get_child_link_name(joint_index)
                control_group.link_names.append(child_name)

        control_group.joint_ids = joint_id_list
        return control_group

    def build_pk_serial_chain(self) -> None:
        """Build the kinematic serial chain for the robot.

        This method is mainly used for robot learning scenarios, for example:
            - Imitation learning dataset generation.
        """
        self.pk_serial_chain = self.cfg.build_pk_serial_chain(device=self.device)

    def set_physical_visible(
        self,
        visible: bool = True,
        control_part: str | None = None,
        rgba: Sequence[float] | None = None,
    ):
        """set collision of the robot or a specific control part.

        Args:
            visible (bool, optional): is collision body visible. Defaults to True.
            control_part (str | None, optional): control part to set visibility. Defaults to None. If None, all links are set.
            rgba (Sequence[float] | None, optional): collision body visible rgba. It will be defined at the first time the function is called. Defaults to None.
        """
        rgba = rgba if rgba is not None else (0.8, 0.2, 0.2, 0.7)
        if len(rgba) != 4:
            logger.log_error(f"Invalid rgba {rgba}, should be a sequence of 4 floats.")
        rgba = np.array(
            [
                rgba[0],
                rgba[1],
                rgba[2],
                rgba[3],
            ]
        )
        link_names = self.get_control_part_link_names(name=control_part)

        # create collision visible node if not exist
        if visible:
            for i, env_idx in enumerate(self._all_indices):
                for link_name in link_names:
                    if self._has_collision_visible_node_dict[link_name] is False:
                        self._entities[env_idx].create_physical_visible_node(
                            rgba, link_name
                        )
                        self._has_collision_visible_node_dict[link_name] = True

        # set visibility
        for i, env_idx in enumerate(self._all_indices):
            for link_name in link_names:
                self._entities[env_idx].set_physical_visible(visible, link_name)

    def destroy(self) -> None:
        return super().destroy()
