# 假设这是在你的 mdp.py 文件中

from isaaclab.envs.mdp.actions.actions_cfg import JointPositionActionCfg
from isaaclab.envs.mdp.actions.joint_actions import JointPositionAction
from isaaclab.utils import configclass
import torch


class CoupledJointPositionAction(JointPositionAction):
    def __init__(self, cfg: 'CoupledJointPositionActionCfg', env): 
        super().__init__(cfg, env)

    @property
    def action_dim(self) -> int:
        """强制 ActionManager 认为只需要 1D 输入。"""
        return 1

    """
    这是运行时被实例化的类。它继承自 JointPositionAction。
    """
    def process_actions(self, actions: torch.Tensor):
        scale = self.cfg.scale
        offset = self.cfg.offset
        # store the raw actions
        self._raw_actions[:] = torch.clamp(actions, -1, 1)
        # apply the affine transformations
        target_position_interval = self._raw_actions * scale + offset
        right_cmd = target_position_interval
        left_cmd = -target_position_interval
        # import pdb; pdb.set_trace()
        self._processed_actions = torch.stack((left_cmd, right_cmd), dim=1).squeeze(-1)
        # print(f"[DEBUG] In: {actions[0]}, {self._processed_actions[0]}")
        # clip actions
        if self.cfg.clip is not None:
            self._processed_actions = torch.clamp(
                self._processed_actions, min=self._clip[:, :, 0], max=self._clip[:, :, 1]
            )


@configclass
class CoupledJointPositionActionCfg(JointPositionActionCfg):
    """
    配置类。关键在于设置 class_type 指向上面的实现类。
    """
    # !!! 关键点 !!! 
    # 告诉 ActionManager: "当你看到这个配置时，请实例化 CoupledJointPositionAction 这个类"
    class_type = CoupledJointPositionAction 

    def __post_init__(self) -> None:
        # 这里的检查逻辑是没问题的，因为 ActionManager 会在初始化时调用它
        if len(self.joint_names) != 2:
            raise ValueError("Requires exactly two joint names.")
        super().__post_init__()

