#!/usr/bin/env python3
"""
验证 teleop_xr_agent.py 中相对模式下动作拼接顺序的正确性。
无需 Isaac Sim / XR 设备，纯 CPU 运行。

运行方式：
    python scripts/debug_action_assembly.py
"""

import torch

# ============================================================
# 模拟 Action Manager 的切分逻辑
# ============================================================
# 来自 mindrobot_left_arm_ik_env_cfg.py 的 Action 配置：
#   arm_action  : DifferentialInverseKinematicsActionCfg，相对模式 6D，绝对模式 7D
#   wheel_action: JointVelocityActionCfg，4 个轮子，4D
#   gripper_action: BinaryJointPositionActionCfg，1D

def action_manager_split(action: torch.Tensor, arm_size: int):
    """模拟 Isaac Lab ManagerBasedRLEnv 的动作切分。"""
    assert action.shape[-1] == arm_size + 4 + 1, \
        f"期望总维度 {arm_size + 4 + 1}，实际 {action.shape[-1]}"
    arm    = action[..., :arm_size]
    wheel  = action[..., arm_size:arm_size + 4]
    gripper = action[..., arm_size + 4:]
    return arm, wheel, gripper


# ============================================================
# 复现 teleop_xr_agent.py:474 的拼接逻辑
# ============================================================

def current_assembly(action_np: torch.Tensor, wheel_np: torch.Tensor) -> torch.Tensor:
    """现有代码的拼接逻辑（逐字复制自 teleop_xr_agent.py:474）。"""
    return torch.cat([action_np[:7], wheel_np, action_np[7:]])


def fixed_assembly(action_np: torch.Tensor, wheel_np: torch.Tensor, is_abs_mode: bool) -> torch.Tensor:
    """修正后的拼接逻辑。"""
    arm_size = 7 if is_abs_mode else 6
    arm_action   = action_np[:arm_size]
    gripper_val  = action_np[arm_size:]   # 最后一个元素
    return torch.cat([arm_action, wheel_np, gripper_val])


# ============================================================
# 测试
# ============================================================

def run_test(mode: str):
    is_abs = (mode == "absolute")
    arm_size = 7 if is_abs else 6

    # 构造可辨识的测试值
    if is_abs:
        # 绝对模式：[x,y,z,qw,qx,qy,qz, gripper]
        fake_action = torch.tensor([0.1, 0.2, 0.3,   # pos
                                    1.0, 0.0, 0.0, 0.0,  # quat
                                    0.9])             # gripper (trigger=0.9)
    else:
        # 相对模式：[dx,dy,dz, drx,dry,drz, gripper]
        fake_action = torch.tensor([0.01, 0.02, 0.03,  # delta pos
                                    0.05, 0.06, 0.07,  # delta rot
                                    0.9])              # gripper (trigger=0.9)

    fake_wheel = torch.tensor([1.5, 1.5, 1.5, 1.5])  # 向前行驶

    print(f"\n{'='*60}")
    print(f"  测试模式: {mode.upper()}")
    print(f"{'='*60}")
    print(f"  原始 action_np  : {fake_action.tolist()}")
    print(f"  原始 wheel_np   : {fake_wheel.tolist()}")
    print(f"  期望 gripper 值 : {fake_action[-1].item()} (应为 0.9)")
    print(f"  期望 wheel  值  : {fake_wheel.tolist()} (应为 [1.5,1.5,1.5,1.5])")

    # ---- 现有代码 ----
    assembled_current = current_assembly(fake_action, fake_wheel)
    try:
        arm_c, wheel_c, grip_c = action_manager_split(assembled_current, arm_size)
        print(f"\n  [现有代码]")
        print(f"    拼接结果({assembled_current.shape[0]}D) : {assembled_current.tolist()}")
        print(f"    → arm     ({arm_c.shape[0]}D) : {arm_c.tolist()}")
        print(f"    → wheel   ({wheel_c.shape[0]}D) : {wheel_c.tolist()}")
        print(f"    → gripper ({grip_c.shape[0]}D) : {grip_c.tolist()}")
        arm_ok    = True
        wheel_ok  = torch.allclose(wheel_c, fake_wheel)
        grip_ok   = torch.allclose(grip_c, fake_action[-1:])
        print(f"    arm 正确? {'✅' if arm_ok else '❌'} | wheel 正确? {'✅' if wheel_ok else '❌  ← BUG'} | gripper 正确? {'✅' if grip_ok else '❌  ← BUG'}")
    except AssertionError as e:
        print(f"    [现有代码] 维度断言失败: {e}")

    # ---- 修正后代码 ----
    assembled_fixed = fixed_assembly(fake_action, fake_wheel, is_abs)
    try:
        arm_f, wheel_f, grip_f = action_manager_split(assembled_fixed, arm_size)
        print(f"\n  [修正后代码]")
        print(f"    拼接结果({assembled_fixed.shape[0]}D) : {assembled_fixed.tolist()}")
        print(f"    → arm     ({arm_f.shape[0]}D) : {arm_f.tolist()}")
        print(f"    → wheel   ({wheel_f.shape[0]}D) : {wheel_f.tolist()}")
        print(f"    → gripper ({grip_f.shape[0]}D) : {grip_f.tolist()}")
        arm_ok    = True
        wheel_ok  = torch.allclose(wheel_f, fake_wheel)
        grip_ok   = torch.allclose(grip_f, fake_action[-1:])
        print(f"    arm 正确? {'✅' if arm_ok else '❌'} | wheel 正确? {'✅' if wheel_ok else '❌'} | gripper 正确? {'✅' if grip_ok else '❌'}")
    except AssertionError as e:
        print(f"    [修正后代码] 维度断言失败: {e}")


if __name__ == "__main__":
    run_test("relative")
    run_test("absolute")
    print()