# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

from __future__ import annotations

import torch
from typing import TYPE_CHECKING

from isaaclab.assets import Articulation, RigidObject
from isaaclab.managers import SceneEntityCfg
from isaaclab.utils.math import wrap_to_pi, quat_error_magnitude, quat_apply, normalize

if TYPE_CHECKING:
    from isaaclab.envs import ManagerBasedRLEnv


def joint_pos_target_l2(env: ManagerBasedRLEnv, target: float, asset_cfg: SceneEntityCfg) -> torch.Tensor:
    """Penalize joint position deviation from a target value."""
    # extract the used quantities (to enable type-hinting)
    asset: Articulation = env.scene[asset_cfg.name]
    # wrap the joint positions to (-pi, pi)
    joint_pos = wrap_to_pi(asset.data.joint_pos[:, asset_cfg.joint_ids])
    # compute the reward
    return torch.sum(torch.square(joint_pos - target), dim=1)


def lid_is_lifted(env: ManagerBasedRLEnv, minimal_height: float,
                  lid_cfg: SceneEntityCfg = SceneEntityCfg("DexCube")) -> torch.Tensor:
    """Reward the agent for lifting the lid above the minimal height."""
    lid: RigidObject = env.scene[lid_cfg.name]
    # root_pos_w shape: (num_envs, 3) - extract z-coordinate: (num_envs,)
    height = lid.data.root_pos_w[:, 2]
    return torch.where(height > minimal_height, 1.0, 0.0)

def gripper_lid_distance(env: ManagerBasedRLEnv, std: float,
                         lid_cfg: SceneEntityCfg = SceneEntityCfg("DexCube"),
                         robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
                         gripper_body_name: str = "left_hand_body") -> torch.Tensor:
    """Reward the agent for reaching the lid using tanh-kernel."""
    # extract the used quantities (to enable type-hinting)
    lid: RigidObject = env.scene[lid_cfg.name]
    robot: Articulation = env.scene[robot_cfg.name]

    # float 化，防止成为 (3,) 向量
    if not isinstance(std, float):
        if torch.is_tensor(std):
            std = std.item()
        else:
            std = float(std)

    # Target lid position: (num_envs, 3)
    lid_pos_w = lid.data.root_pos_w[:, :3]
    
    # Gripper position: (num_envs, 3)
    # body_idx = robot.find_bodies(gripper_body_name)[0]
    # gripper_pos_w = robot.data.body_pos_w[:, body_idx, :].squeeze(dim=1)
    body_ids, _ = robot.find_bodies([gripper_body_name], preserve_order=True)
    gripper_pos_w = robot.data.body_pos_w[:, body_ids[0], :]
    # Distance of the gripper to the lid: (num_envs,)
    distance = torch.norm(lid_pos_w - gripper_pos_w, dim=1) + 1e-6
    
    return 1 - torch.tanh(distance / std)


def lid_grasped(env: ManagerBasedRLEnv,
                distance_threshold: float = 0.05,
                lid_cfg: SceneEntityCfg = SceneEntityCfg("DexCube"),
                robot_cfg: SceneEntityCfg = SceneEntityCfg     ("Mindbot"),
                gripper_body_name: str = "left_hand_body") -> torch.Tensor:
    """Reward the agent for grasping the lid (close and lifted)."""
    # extract the used quantities (to enable type-hinting)
    lid: RigidObject = env.scene[lid_cfg.name]
    robot: Articulation = env.scene[robot_cfg.name]
    
    # Target lid position: (num_envs, 3)
    lid_pos_w = lid.data.root_pos_w[:, :3]
    
    # Gripper position: (num_envs, 3)
    body_idx = robot.find_bodies(gripper_body_name)[0]
    gripper_pos_w = robot.data.body_pos_w[:, body_idx, :]
    
    # Distance of the gripper to the lid: (num_envs,)
    distance = torch.norm(lid_pos_w - gripper_pos_w, dim=1)
    
    # Check if close and lifted
    is_close = distance < distance_threshold
    is_lifted = lid.data.root_pos_w[:, 2] > 0.1
    
    return torch.where(is_close & is_lifted, 1.0, 0.0)

def gripper_lid_distance_combined(env: ManagerBasedRLEnv, std: float, minimal_height: float,
                                  lid_cfg: SceneEntityCfg = SceneEntityCfg("DexCube"),
                                  robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
                                  gripper_body_name: str = "left_hand_body") -> torch.Tensor:
    """Combined reward for reaching the lid AND lifting it."""
    # Get reaching reward
    reach_reward = gripper_lid_distance(env, std, lid_cfg, robot_cfg, gripper_body_name)
    # Get lifting reward
    lift_reward = lid_is_lifted(env, minimal_height, lid_cfg)
    # Combine rewards multiplicatively
    return reach_reward * lift_reward


def gripper_handle_orientation_alignment(
    env,
    dryingbox_cfg: SceneEntityCfg = SceneEntityCfg("dryingbox"),
    robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
    handle_name: str = "Equipment_BB_13_C",
    gripper_body_name: str = "right_hand_body",
    gripper_forward_axis: tuple = (0.0, 0.0, 1.0),  # 夹爪指尖指向 (Local Z)
    gripper_width_axis: tuple = (0.0, 1.0, 0.0),    # 夹爪手指开合轴 (Local Y)
    lid_handle_axis: tuple = (0.0, 0.0, 1.0),  
    max_angle_deg: float =60.0,
    # handle_local_axis 参数已删除，因为直接使用世界坐标系 Z 轴更稳健
) -> torch.Tensor:
    """
    修正版：
    1. 解决了梯度消失：使用 (dot+1)/2 替代 clamp，保证全角度有梯度。
    2. 解决了轴向错误：直接认定手柄是垂直的 (World Z)，修复抓取姿态。
    """
    # 1. 获取数据
    # dryingbox = env.scene[dryingbox_cfg.name] (计算轴向不再需要它，节省计算)
    robot = env.scene[robot_cfg.name]
    body_ids, _ = robot.find_bodies([gripper_body_name], preserve_order=True)
    gripper_quat_w = robot.data.body_quat_w[:, body_ids[0], :]

    box=env.scene[dryingbox_cfg.name]
    hd_ids,_=box.find_bodies([handle_name],preserve_order=True)
    handle_quat_w=box.data.body_quat_w[:,hd_ids[0],:]

    # --- 约束 A: 指向约束 (Pointing) ---
    # 目标：夹爪 Z 轴 指向 World +X
    grip_fwd_local = torch.tensor(gripper_forward_axis, device=env.device).repeat(env.num_envs, 1)
    grip_fwd_world = quat_apply(gripper_quat_w, grip_fwd_local)
    
    target_fwd = torch.tensor([1.0, 0.0, 0.0], device=env.device).repeat(env.num_envs, 1)
    
    dot_fwd = torch.sum(grip_fwd_world * target_fwd, dim=1)
    dot_fwd=torch.clamp(dot_fwd,-1.0,1.0)
    angle_rad = torch.acos(torch.clamp(dot_fwd, -1.0, 1.0))
    max_angle_rad = torch.tensor(max_angle_deg * 3.14159265359 / 180.0, device=env.device)
    angle_normalized = torch.clamp(angle_rad / max_angle_rad, 0.0, 1.0)
    rew_forward = 1.0 - angle_normalized  # 角度越小，奖励越大
    # 如果角度超过60度，奖励为0
    rew_forward = torch.where(angle_rad <= max_angle_rad, rew_forward, torch.zeros_like(rew_forward))
    # nan -> 0
    rew_forward = torch.where(torch.isnan(rew_forward), torch.zeros_like(rew_forward), rew_forward)

    # [关键修改 1] 使用软化奖励，提供全局梯度
    # 范围 [0, 1]: 1.0 是完美对齐，0.5 是垂直，0.0 是完全背对
    # 这样即使现在是指向地面的 (dot=0)，它也有 0.5 分，且知道往上抬分会变高
    # rew_forward = (dot_fwd + 1.0) / 2.0
    # 如果你希望能更陡峭一点，可以加个平方： rew_forward = rew_forward * rew_forward

    # --- 约束 B: 抓取姿态约束 (Grasp Alignment) ---
    # 目标：夹爪的开合轴 (Y) 应该垂直于 手柄轴 (Z)
    # 这样手指就是水平开合的
    
    grip_width_local = torch.tensor(gripper_width_axis, device=env.device).repeat(env.num_envs, 1)
    grip_width_world = normalize(quat_apply(gripper_quat_w, grip_width_local))
    
    # [关键修改 2] 强制假设手柄是垂直的 (World Z)
    # 除非你的烘箱会倾斜倒在地上，否则这是最稳健的写法，不用管局部坐标系
    handle_axis_local = torch.tensor([0.0, 0.0, 1.0], device=env.device).repeat(env.num_envs, 1)
    handle_axis_world=quat_apply(handle_quat_w,handle_axis_local)
    dot_yaw=torch.sum(grip_width_world*handle_axis_world,dim=1)
    rew_yaw=torch.square(torch.abs(dot_yaw))+1e-4
    # nan -> 0
    rew_yaw = torch.where(torch.isnan(rew_yaw), torch.zeros_like(rew_yaw), rew_yaw)

    # 5. 组合
    total_reward = rew_forward * rew_yaw

    total_reward = torch.nan_to_num(total_reward, nan=0.0, posinf=0.0, neginf=0.0)
    # debug
    if not torch.isfinite(total_reward).all():
        print(f"[DEBUG] Non-finite reward detected in gripper_lid_orientation_alignment: {total_reward}")
        print(f"rew_vertical: {rew_vertical}")
        print(f"rew_yaw: {rew_yaw}")
        
    return total_reward

    # # 计算点积：我们希望它们垂直，即点积为 0
    # dot_width_handle = torch.sum(grip_width_world * handle_axis_world, dim=1)
    
    # # 奖励：1.0 - abs(dot)。点积越小(越垂直)，分数越高
    # rew_perpendicular = 1.0 - torch.abs(dot_width_handle)
    
    # # 稍微增强一下敏感度，让它更倾向于完美垂直
    # rew_perpendicular = rew_perpendicular * rew_perpendicular

    # # 3. 组合奖励
    # total_reward = rew_forward * rew_perpendicular
    
    # return torch.nan_to_num(total_reward, nan=0.0)


# def gripper_handle_position_alignment(env: ManagerBasedRLEnv,
#                                     dryingbox_cfg: SceneEntityCfg = SceneEntityCfg("dryingbox"),
#                                     robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
#                                     handle_name: str = "Equipment_BB_13_C",
#                                     left_gripper_name: str = "left_hand__l",
#                                     right_gripper_name: str = "left_hand_r",
#                                     height_offset: float = 0.03,
#                                     std: float = 0.1) -> torch.Tensor:

#     dryingbox: RigidObject = env.scene[dryingbox_cfg.name]
#     robot: Articulation = env.scene[robot_cfg.name]
    
#     # lid_pos_w = lid.data.root_pos_w[:, :3]
#     handle_ids, _ = dryingbox.find_bodies([handle_name], preserve_order=True)
#     # [num_envs, 3]
#     handle_pos_w = dryingbox.data.body_pos_w[:, handle_ids[0], :]
#     # 目标点：lid位置 + height_offset（Z方向）
#     target_pos = handle_pos_w.clone()
#     # print(target_pos[0])
#     target_pos[:, 0] -= height_offset  # 在lid上方height_offset处
#     # print(target_pos[0])
#     # 获取两个夹爪的位置
#     left_body_ids, _ = robot.find_bodies([left_gripper_name], preserve_order=True)
#     right_body_ids, _ = robot.find_bodies([right_gripper_name], preserve_order=True)
    
#     left_gripper_pos = robot.data.body_pos_w[:, left_body_ids[0], :]  # (num_envs, 3)
#     right_gripper_pos = robot.data.body_pos_w[:, right_body_ids[0], :]  # (num_envs, 3)
    
#     # 计算夹爪中心位置
#     gripper_center = (left_gripper_pos + right_gripper_pos) / 2.0  # (num_envs, 3)
#     # print(gripper_center[0])
#     # 计算夹爪中心到目标点的3D距离
#     distance = torch.norm(gripper_center - target_pos, dim=1) + 1e-6
    
#     # 距离奖励：距离越小，奖励越大
#     position_reward = 1.0 - torch.tanh(distance / std)
    
#     return position_reward


# def gripper_handle_position_alignment(
#     env,
#     dryingbox_cfg: SceneEntityCfg = SceneEntityCfg("dryingbox"),
#     robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
#     handle_name: str = "Equipment_BB_13_C",
#     left_finger_name: str = "right_hand_l",
#     right_finger_name: str = "right_hand__r",
#     # 【修改点1】：根据你的描述，长轴是 Local Z，所以改为 (0, 0, 1)
#     handle_local_axis: tuple = (0.0, 0.0, 1.0), 
#     # 【新增】：明确指定手柄的“侧面”是 Local X 轴
#     handle_local_side_axis: tuple = (0.0, -1.0, 0.0),
#     grasp_width: float = 0.03,
#     grasp_offset: float = 0.06,
#     std: float = 0.05,
# ) -> torch.Tensor:
#     """
#     修正版：
#     1. 修正了手柄长轴为 Local Z。
#     2. 使用 Local X 作为手柄两侧的基准方向。
#     """
#     # 1. 获取数据
#     dryingbox = env.scene[dryingbox_cfg.name]
#     robot = env.scene[robot_cfg.name]

#     handle_ids, _ = dryingbox.find_bodies([handle_name], preserve_order=True)

#     # [num_envs, 3]
#     handle_pos_w = dryingbox.data.body_pos_w[:, handle_ids[0], :]
#     # [num_envs, 4]
#     handle_quat_w = dryingbox.data.body_quat_w[:, handle_ids[0], :]

#     l_finger_ids, _ = robot.find_bodies([left_finger_name], preserve_order=True)
#     r_finger_ids, _ = robot.find_bodies([right_finger_name], preserve_order=True)
#     l_finger_pos = robot.data.body_pos_w[:, l_finger_ids[0], :]
#     r_finger_pos = robot.data.body_pos_w[:, r_finger_ids[0], :]

#     # 2. 计算关键向量
#     # 【修正】：计算世界坐标系下的手柄长轴 (应为 World Z)
#     handle_axis_local_vec = torch.tensor(handle_local_axis, device=env.device).repeat(env.num_envs, 1)
#     # 虽然本奖励函数主要位置点，不需要显式用到长轴向量做点积，但保留以备后续增加角度奖励
#     handle_axis_world = normalize(quat_apply(handle_quat_w, handle_axis_local_vec)) 

#     # 【核心修改】：直接计算世界坐标系下的手柄“侧向”向量 (即 Local X -> World Y)
#     # 这种方法比 cross(approach, axis) 更准确，因为它尊重了物体自身的旋转
#     handle_side_local_vec = torch.tensor(handle_local_side_axis, device=env.device).repeat(env.num_envs, 1)
#     side_vec_world = normalize(quat_apply(handle_quat_w, handle_side_local_vec))

#     # 假设抓取方向主要是沿世界 X 轴 (指向烘箱)
#     # 注意：如果你的烘箱放置位置旋转了，这里也建议改为读取机器人的前方或者物体的Local方向
#     approach_dir = torch.tensor([1.0, 0.0, 0.0], device=env.device).repeat(env.num_envs, 1)

#     # 3. 计算理想的抓取点
#     # 抓取中心：从手柄中心沿着抓取方向（World X）反向退后 grasp_offset
#     grasp_center = handle_pos_w - approach_dir * grasp_offset

#     # 计算理想的左/右手指位置
#     # 既然 side_vec_world 是 Local X (对应 World Y)
#     # 我们希望手指在 Y 轴两侧张开
#     ideal_l_pos = grasp_center + side_vec_world * (grasp_width / 2.0)
#     ideal_r_pos = grasp_center - side_vec_world * (grasp_width / 2.0)
    
#     # 注意：这里 l_pos 用 + 还是 - 取决于你的 robot 建模里 right_hand_l 实际是在正Y侧还是负Y侧
#     # 如果发现左右手指交叉，交换上面的 + 和 - 即可。

#     # 4. 计算距离奖励
#     dist_l = torch.norm(l_finger_pos - ideal_r_pos, dim=1)
#     dist_r = torch.norm(r_finger_pos - ideal_l_pos, dim=1)

#     rew_l = 1.0 - torch.tanh(dist_l / std)
#     rew_r = 1.0 - torch.tanh(dist_r / std)

#     total_reward = (rew_l + rew_r) / 2.0

#     return torch.nan_to_num(total_reward, nan=0.0)

# def gripper_handle_position_alignment(
#     env,
#     dryingbox_cfg: SceneEntityCfg = SceneEntityCfg("dryingbox"),
#     robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
#     handle_name: str = "Equipment_BB_13_C",
#     left_finger_name: str = "right_hand_l",     # 左手指 body 名称
#     right_finger_name: str = "right_hand__r",   # 右手指 body 名称
#     handle_local_axis: tuple = (0.0, 0.0, 1.0), # 手柄长轴
#     grasp_width: float = 0.03,                  # 预期的抓取宽度 (米)
#     grasp_offset: float = 0.03,                 # 手指尖距离手柄中心的进深偏移
#     std: float = 0.05,                          # 奖励函数的敏感度
# ) -> torch.Tensor:
#     """
#     双指位置奖励：引导左/右手指分别到达手柄的两侧。
#     """
#     # 1. 获取数据
#     dryingbox = env.scene[dryingbox_cfg.name]
#     robot = env.scene[robot_cfg.name]

#     handle_ids, _ = dryingbox.find_bodies([handle_name], preserve_order=True)
#     handle_pos_w = dryingbox.data.body_pos_w[:, handle_ids[0], :]
#     handle_quat_w = dryingbox.data.body_quat_w[:, handle_ids[0], :]

#     l_finger_ids, _ = robot.find_bodies([left_finger_name], preserve_order=True)
#     r_finger_ids, _ = robot.find_bodies([right_finger_name], preserve_order=True)
#     l_finger_pos = robot.data.body_pos_w[:, l_finger_ids[0], :]
#     r_finger_pos = robot.data.body_pos_w[:, r_finger_ids[0], :]

#     # 2. 计算理想的抓取点
#     # 获取手柄当前轴向
#     handle_axis_local_vec = torch.tensor(handle_local_axis, device=env.device).repeat(env.num_envs, 1)
#     handle_axis_world = normalize(quat_apply(handle_quat_w, handle_axis_local_vec))

#     # 假设抓取方向主要是沿世界 X 轴 (指向烘箱)
#     approach_dir = torch.tensor([1.0, 0.0, 0.0], device=env.device).repeat(env.num_envs, 1)

#     # 计算垂直于手柄轴向和接近方向的向量，这就是“侧面”方向
#     side_vec = normalize(torch.cross(approach_dir, handle_axis_world))

#     # 计算抓取中心点 (手柄中心沿接近方向反向偏移一点，留出空间)
#     grasp_center = handle_pos_w - approach_dir * grasp_offset

#     # 计算理想的左/右手指位置 (在中心点两侧张开)
#     # 注意：这里需要根据你的坐标系确认哪个是左，哪个是右
#     ideal_l_pos = grasp_center - side_vec * (grasp_width / 2.0)
#     ideal_r_pos = grasp_center + side_vec * (grasp_width / 2.0)

#     # 3. 计算距离奖励
#     dist_l = torch.norm(l_finger_pos - ideal_l_pos, dim=1)
#     dist_r = torch.norm(r_finger_pos - ideal_r_pos, dim=1)

#     rew_l = 1.0 - torch.tanh(dist_l / std)
#     rew_r = 1.0 - torch.tanh(dist_r / std)

#     # 取平均值
#     total_reward = (rew_l + rew_r) / 2.0

#     return torch.nan_to_num(total_reward, nan=0.0)


# 在 rewards.py 中，修改 _is_grasped 函数，使其支持 dryingbox 和 handle_name
def _is_grasped(
    env: ManagerBasedRLEnv,
    dryingbox_cfg: SceneEntityCfg,  # 改为 dryingbox_cfg
    robot_cfg: SceneEntityCfg,
    handle_name: str,  # 新增 handle_name 参数
    left_gripper_name: str,
    right_gripper_name: str,
    joint_names: list,
    height_offset: float,
    grasp_radius: float,
    target_close_pos: float
) -> torch.Tensor:
    """
    内部辅助函数：判定是否成功抓取。
    条件：
    1. 夹爪中心在把手目标点附近 (Sphere Check)
    2. 夹爪处于闭合发力状态 (Joint Effort Check)
    3. (关键) 夹爪X轴距离合适，不能压在把手上 (X-Axis Check)
    """
    # 1. 获取对象
    dryingbox = env.scene[dryingbox_cfg.name]  # 改为 dryingbox
    robot = env.scene[robot_cfg.name]
    
    # 2. 目标点位置 (把手中心)
    handle_ids, _ = dryingbox.find_bodies([handle_name], preserve_order=True)
    handle_pos_w = dryingbox.data.body_pos_w[:, handle_ids[0], :3].clone()
    handle_pos_w[:, 1] += height_offset  # X方向偏移
    handle_pos_w[:, 0] += 0.02
    # handle_pos_w[:, 2] -= 0.04
    # 3. 夹爪位置
    l_ids, _ = robot.find_bodies([left_gripper_name])
    r_ids, _ = robot.find_bodies([right_gripper_name])
    pos_L = robot.data.body_pos_w[:, l_ids[0], :3]
    pos_R = robot.data.body_pos_w[:, r_ids[0], :3]
    gripper_center = (pos_L + pos_R) / 2.0
    
    # 4. 条件一：距离判定 (在把手球范围内)
    dist_to_handle = torch.norm(gripper_center - handle_pos_w, dim=-1)
    if env.common_step_counter % 100 == 0:
        print(f"Env 0 Dist to Handle: {dist_to_handle[0].item():.4f}m")
    is_near = dist_to_handle < grasp_radius
    
    # 5. 条件二：X轴距离判定 (防止压在把手上)
    x_diff = gripper_center[:, 0] - handle_pos_w[:, 0]
    is_x_valid = torch.abs(x_diff) < 0.03  # 3cm 容差
    
    # 6. 条件三：闭合判定 (正在尝试闭合)
    joint_ids, _ = robot.find_joints(joint_names)
    joint_pos = torch.abs(robot.data.joint_pos[:, joint_ids])
    
    is_effort_closing = torch.all(joint_pos > 0.005, dim=1)
    
    # 2. 几何排斥判定 (Geometry Check)
    dist_fingertips = torch.norm(pos_L - pos_R, dim=-1)
    is_not_empty = dist_fingertips > 0.005

    # 在计算 dist_to_handle 后
    if env.common_step_counter % 500 == 0 and dist_to_handle[0] < 0.15:
        delta = gripper_center[0] - handle_pos_w[0]
        print(f"close_dbg env0: handle={handle_pos_w[0].cpu().numpy()}, "
                f"center={gripper_center[0].cpu().numpy()}, "
                f"delta={delta.cpu().numpy()}, "
                f"dist={dist_to_handle[0].item():.4f}, "
                f"fingertip_dist={dist_fingertips[0].item():.4f}, "
                f"joint_pos={joint_pos[0].cpu().numpy()}")
        
    # 综合判定
    return is_near & is_x_valid & is_effort_closing & is_not_empty


# 修改 gripper_close_when_near 函数
def gripper_close_when_near(
    env: ManagerBasedRLEnv,
    dryingbox_cfg: SceneEntityCfg = SceneEntityCfg("dryingbox"),  # 改为 dryingbox_cfg
    robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
    handle_name: str = "Equipment_BB_13_C",  # 新增 handle_name 参数
    left_gripper_name: str = "right_hand_l",  # 更新默认值
    right_gripper_name: str = "right_hand__r",  # 更新默认值
    joint_names: list = ["right_hand_joint_left", "right_hand_joint_right"],  # 更新默认值
    target_close_pos: float = 0.03,
    height_offset: float = 0.005,  # 更新默认值
    grasp_radius: float = 0.03  # 更新默认值
) -> torch.Tensor:
    
    # 1. 判定基础抓取条件是否满足
    is_grasped = _is_grasped(
        env, dryingbox_cfg, robot_cfg, handle_name,  # 传入 handle_name
        left_gripper_name, right_gripper_name,
        joint_names, height_offset, grasp_radius, target_close_pos
    )
    
    # 2. 计算夹紧程度 (Clamping Intensity)
    robot = env.scene[robot_cfg.name]
    joint_ids, _ = robot.find_joints(joint_names)
    current_joint_pos = torch.abs(robot.data.joint_pos[:, joint_ids])
    mean_joint_pos = torch.mean(current_joint_pos, dim=1)
    
    clamping_factor = torch.clamp(mean_joint_pos / target_close_pos, max=1.0)
    # 3. 只有在满足抓取判定时，才发放带有权重的夹紧奖励



    return torch.where(is_grasped, clamping_factor, 0.0)



# def gripper_close_when_near(
#     env: ManagerBasedRLEnv,
#     lid_cfg: SceneEntityCfg = SceneEntityCfg("lid"),
#     robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
#     # 注意：这里我们需要具体的指尖 body 名字来做几何判定
#     left_gripper_name: str = "left_hand__l", # 请确认你的USD里左指尖body名
#     right_gripper_name: str = "left_hand_r", # 请确认你的USD里右指尖body名
#     joint_names: list = ["left_hand_joint_left", "left_hand_joint_right"], 
#     target_close_pos: float = 0.03,
#     height_offset: float = 0.03,
#     grasp_radius: float = 0.02 # 球面半径 2cm
# ) -> torch.Tensor:
    
#     # 1. 获取位置
#     lid = env.scene[lid_cfg.name]
#     robot = env.scene[robot_cfg.name]
#     lid_pos_w = lid.data.root_pos_w[:, :3].clone()
#     lid_pos_w[:, 2] += height_offset # 把手中心
    
#     # 获取左右指尖位置
#     l_ids, _ = robot.find_bodies([left_gripper_name])
#     r_ids, _ = robot.find_bodies([right_gripper_name])
#     pos_L = robot.data.body_pos_w[:, l_ids[0], :3]
#     pos_R = robot.data.body_pos_w[:, r_ids[0], :3]
    
#     # 计算夹爪中心
#     pos_center = (pos_L + pos_R) / 2.0
    
#     # 2. 距离判定 (Is Inside Sphere?)
#     dist_center = torch.norm(pos_center - lid_pos_w, dim=-1)
#     is_in_sphere = (dist_center < grasp_radius).float()
    
#     # 3. "在中间"判定 (Is Between?)
#     # 投影逻辑：物体投影在 LR 连线上，且位于 L 和 R 之间
#     vec_LR = pos_R - pos_L # 左指指向右指
#     vec_LO = lid_pos_w - pos_L # 左指指向物体
    
#     # 计算投影比例 t
#     # P_proj = P_L + t * (P_R - P_L)
#     # t = (vec_LO . vec_LR) / |vec_LR|^2
#     # 如果 0 < t < 1，说明投影在两个指尖之间
    
#     len_sq = torch.sum(vec_LR * vec_LR, dim=-1) + 1e-6
#     dot = torch.sum(vec_LO * vec_LR, dim=-1)
#     t = dot / len_sq
    
#     is_between = (t > 0.0) & (t < 1.0)
#     is_between = is_between.float()
    
#     # 4. 闭合判定
#     joint_ids, _ = robot.find_joints(joint_names)
#     # 注意：如果是 Binary Action，可能很难拿到连续的 0~0.03 控制值，如果是仿真状态则没问题
#     current_joint_pos = torch.abs(robot.data.joint_pos[:, joint_ids])
#     close_error = torch.mean(torch.abs(current_joint_pos - target_close_pos), dim=1)
#     # 只有当非常接近闭合目标时才给分
#     is_closing = (close_error < 0.005).float() # 允许 5mm 误差
    
#     # 5. 最终奖励
#     # 只有三者同时满足才给 1.0 分
#     reward = is_in_sphere * is_between * is_closing
    
#     return torch.nan_to_num(reward, nan=0.0)


# def lid_is_lifted(
#     env:ManagerBasedRLEnv,
#     minimal_height:float,
#     lid_cfg:SceneEntityCfg = SceneEntityCfg("lid")
# ) -> torch.Tensor:

#     lid = env.scene[lid_cfg.name]
#     lid_height = lid.data.root_pos_w[:, 2]
#     reward=torch.where(lid_height > minimal_height, 1.0, 0.0)
#     reward=torch.nan_to_num(reward, nan=0.0)
    
#     return reward


def lid_is_lifted(
    env: ManagerBasedRLEnv,
    minimal_height: float = 0.05, # 相对提升阈值
    lid_cfg: SceneEntityCfg = SceneEntityCfg("lid"),
    robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
    left_gripper_name: str = "left_hand__l",
    right_gripper_name: str = "left_hand_r",
    joint_names: list = ["left_hand_joint_left", "left_hand_joint_right"], 
    height_offset: float = 0.03,
    grasp_radius: float = 0.1,
    target_close_pos: float = 0.03,
) -> torch.Tensor:
    
    is_grasped = _is_grasped(
        env, lid_cfg, robot_cfg, left_gripper_name, right_gripper_name,
        joint_names, height_offset, grasp_radius, target_close_pos
    )

    lid = env.scene[lid_cfg.name]
    
    # 1. 获取高度
    current_height = lid.data.root_pos_w[:, 2]
    # 自动获取初始高度
    initial_height = lid.data.default_root_state[:, 2]
    
    # 2. 计算提升量
    lift_height = torch.clamp(current_height - initial_height, min=0.0)
    
    # 3. 速度检查 (防撞飞)
    # 如果速度 > 1.0 m/s，视为被撞飞，不给分
    lid_vel = torch.norm(lid.data.root_lin_vel_w, dim=1)
    is_stable = lid_vel < 1.0
    
    # 4. 计算奖励
    # 基础分：高度越高分越高
    shaping_reward = lift_height * 2.0
    # 成功分：超过阈值
    success_bonus = torch.where(lift_height > minimal_height, 1.0, 0.0)
    
    # 组合
    # 必须 is_grasped AND is_stable
    reward = torch.where(is_stable & is_grasped, shaping_reward + success_bonus, torch.zeros_like(lift_height))
    
    return torch.nan_to_num(reward, nan=0.0)



def lid_lift_success_reward(
    env: ManagerBasedRLEnv,
    lid_cfg: SceneEntityCfg = SceneEntityCfg("lid"),
    robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
    left_gripper_name: str = "left_hand_body",
    joint_names: list = ["left_hand_joint_left", "left_hand_joint_right"],
    # 关键参数
    settled_height: float = 0.75,  # 盖子落在超声仪上的稳定高度 (需根据仿真实际情况微调)
    target_lift_height: float = 0.05, # 目标提升高度 (5cm)
    grasp_dist_threshold: float = 0.05, # 抓取判定距离
    closed_pos: float = 0.03 # 夹爪闭合阈值
) -> torch.Tensor:
    """
    提升奖励：只有在【夹稳】的前提下，将盖子【相对】提升达到目标高度才给高分。
    """
    # 1. 获取数据
    lid = env.scene[lid_cfg.name]
    robot = env.scene[robot_cfg.name]
    
    # 盖子实时高度
    lid_z = lid.data.root_pos_w[:, 2]
    lid_pos = lid.data.root_pos_w[:, :3]
    
    # 夹爪位置
    body_ids, _ = robot.find_bodies([left_gripper_name])
    gripper_pos = robot.data.body_pos_w[:, body_ids[0], :3]
    
    # 2. 【条件一：是否夹稳 (Is Grasped?)】
    # 距离检查
    distance = torch.norm(gripper_pos - lid_pos, dim=-1)
    is_close = distance < grasp_dist_threshold
    
    # 闭合检查
    joint_ids, _ = robot.find_joints(joint_names)
    joint_pos_abs = torch.abs(robot.data.joint_pos[:, joint_ids])
    # 假设 > 80% 的闭合度算抓紧了
    is_closed = torch.all(joint_pos_abs > (closed_pos * 0.8), dim=-1)
    
    is_grasped = is_close & is_closed
    
    # 3. 【条件二：计算相对提升 (Relative Lift)】
    # 当前高度 - 初始稳定高度
    current_lift = lid_z - settled_height
    
    # 4. 计算奖励
    lift_progress = torch.clamp(current_lift / target_lift_height, min=0.0, max=1.0)
    
    # 基础提升分 (Shaping Reward)
    lift_reward = lift_progress
    
    # 成功大奖 (Success Bonus)
    # 如果提升超过目标高度 (比如 > 95% 的 5cm)，给额外大奖
    success_bonus = torch.where(current_lift >= target_lift_height, 2.0, 0.0)
    
    # 组合：只有在 is_grasped 为 True 时才发放提升奖励
    total_reward = torch.where(is_grasped, lift_reward + success_bonus, torch.zeros_like(lift_reward))
    
    # 5. 安全输出
    return torch.nan_to_num(total_reward, nan=0.0)

def lid_lift_progress_reward(
    env: ManagerBasedRLEnv,
    lid_cfg: SceneEntityCfg = SceneEntityCfg("lid"),
    robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
    left_gripper_name: str = "left_hand__l",
    right_gripper_name: str = "left_hand_r",
    joint_names: list = ["left_hand_joint_left", "left_hand_joint_right"], 
    initial_height: float = 0.8, 
    target_height_lift: float = 0.15,
    height_offset: float = 0.02, # 必须与抓取奖励保持一致
    grasp_radius: float = 0.1,   # 提升时可以稍微放宽一点半径，防止抖动丢分
    target_close_pos: float = 0.03,
    std: float = 0.05 # 标准差
) -> torch.Tensor:

    # 1. 判定是否夹住
    is_grasped = _is_grasped(
        env, lid_cfg, robot_cfg, left_gripper_name, right_gripper_name,
        joint_names, height_offset, grasp_radius, target_close_pos
    )
    
    # 2. 计算高度
    lid = env.scene[lid_cfg.name]
    current_height = lid.data.root_pos_w[:, 2]
    lift_height = torch.clamp(current_height - initial_height, min=0.0, max=target_height_lift)
    # print(current_height)
    # print(lift_height)
    # 3. 计算奖励
    # 只有在 is_grasped 为 True 时，才发放高度奖励
    # 这样彻底杜绝了 "砸飞/撞飞" 拿分的情况
    progress = torch.tanh(lift_height / std)
    
    reward = torch.where(is_grasped, progress, 0.0)
    
    return reward

# def lid_grasped_bonus(
#     env: ManagerBasedRLEnv,
#     lid_cfg: SceneEntityCfg,
#     robot_cfg: SceneEntityCfg,
#     left_gripper_name: str,
#     joint_names: list,
#     distance_threshold: float = 0.05,
#     closed_pos: float = 0.03
# ) -> torch.Tensor:
#     # 1. 计算距离
#     lid = env.scene[lid_cfg.name]
#     robot = env.scene[robot_cfg.name]
#     body_ids, _ = robot.find_bodies([left_gripper_name])
#     gripper_pos = robot.data.body_pos_w[:, body_ids[0], :3]
#     lid_pos = lid.data.root_pos_w[:, :3]
#     distance = torch.norm(gripper_pos - lid_pos, dim=-1)
    
#     # 2. 检查距离是否达标
#     is_close = distance < distance_threshold
    
#     # 3. 检查夹爪是否闭合
#     joint_ids, _ = robot.find_joints(joint_names)
#     joint_pos = torch.abs(robot.data.joint_pos[:, joint_ids])
#     # 检查是否接近闭合 (比如 > 0.02)
#     is_closed = torch.all(joint_pos > (closed_pos * 0.8), dim=-1)
    
#     # 4. 给分
#     # 只有同时满足才给 1.0
#     reward = torch.where(is_close & is_closed, 1.0, 0.0)
    
#     return torch.nan_to_num(reward)

# 在 rewards.py 文件开头添加调试函数
def debug_robot_physics(
    env: ManagerBasedRLEnv,
    robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
    gripper_body_name: str = "right_hand_body",
    distance_threshold: float = 0.1,  # 当距离小于此值时打印
    print_interval: int = 10  # 每N步打印一次
) -> None:
    """
    调试函数：打印机器人的速度、加速度等信息
    """
    # 使用静态变量控制打印频率
    if not hasattr(debug_robot_physics, "step_count"):
        debug_robot_physics.step_count = 0
    debug_robot_physics.step_count += 1
    
    if debug_robot_physics.step_count % print_interval != 0:
        return
    
    robot = env.scene[robot_cfg.name]
    
    # 1. 获取机器人根节点的速度和加速度
    root_lin_vel = robot.data.root_lin_vel_w  # (num_envs, 3)
    root_ang_vel = robot.data.root_ang_vel_w  # (num_envs, 3)
    root_lin_vel_magnitude = torch.norm(root_lin_vel, dim=1)  # (num_envs,)
    root_ang_vel_magnitude = torch.norm(root_ang_vel, dim=1)  # (num_envs,)
    
    # 2. 获取夹爪body的速度
    body_ids, _ = robot.find_bodies([gripper_body_name], preserve_order=True)
    gripper_lin_vel = robot.data.body_lin_vel_w[:, body_ids[0], :]  # (num_envs, 3)
    gripper_ang_vel = robot.data.body_ang_vel_w[:, body_ids[0], :]  # (num_envs, 3)
    gripper_lin_vel_mag = torch.norm(gripper_lin_vel, dim=1)
    gripper_ang_vel_mag = torch.norm(gripper_ang_vel, dim=1)
    
    # 3. 获取关节速度（可能影响碰撞）
    joint_vel = robot.data.joint_vel  # (num_envs, num_joints)
    max_joint_vel = torch.max(torch.abs(joint_vel), dim=1)[0]  # (num_envs,)
    
    # 4. 检查是否有异常高的速度（可能导致碰撞问题）
    high_vel_mask = (gripper_lin_vel_mag > 0.5) | (root_lin_vel_magnitude > 0.3)
    
    if torch.any(high_vel_mask):
        env_ids = torch.where(high_vel_mask)[0].cpu().tolist()
        for env_id in env_ids[:3]:  # 只打印前3个环境
            print(f"\n[DEBUG] Step {debug_robot_physics.step_count}, Env {env_id}:")
            print(f"  机器人根节点线性速度: {root_lin_vel[env_id].cpu().numpy()} (magnitude: {root_lin_vel_magnitude[env_id].item():.4f} m/s)")
            print(f"  机器人根节点角速度: {root_ang_vel[env_id].cpu().numpy()} (magnitude: {root_ang_vel_magnitude[env_id].item():.4f} rad/s)")
            print(f"  夹爪线性速度: {gripper_lin_vel[env_id].cpu().numpy()} (magnitude: {gripper_lin_vel_mag[env_id].item():.4f} m/s)")
            print(f"  夹爪角速度: {gripper_ang_vel[env_id].cpu().numpy()} (magnitude: {gripper_ang_vel_mag[env_id].item():.4f} rad/s)")
            print(f"  最大关节速度: {max_joint_vel[env_id].item():.4f} rad/s")

def gripper_handle_position_alignment(env: ManagerBasedRLEnv,
                                    dryingbox_cfg: SceneEntityCfg = SceneEntityCfg("dryingbox"),
                                    robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
                                    handle_name: str = "Equipment_BB_13_C",
                                    left_gripper_name: str = "right_hand_l",
                                    right_gripper_name: str = "right_hand__r",
                                    height_offset: float = 0.03,
                                    std: float = 0.1) -> torch.Tensor:

    dryingbox: RigidObject = env.scene[dryingbox_cfg.name]
    robot: Articulation = env.scene[robot_cfg.name]
    
    handle_ids, _ = dryingbox.find_bodies([handle_name], preserve_order=True)
    handle_pos_w = dryingbox.data.body_pos_w[:, handle_ids[0], :]
    target_pos = handle_pos_w.clone()
    target_pos[:, 1] += height_offset
    target_pos[:, 0] += 0.02
    # target_pos[:, 2] -= 0.04
    
    left_body_ids, _ = robot.find_bodies([left_gripper_name], preserve_order=True)
    right_body_ids, _ = robot.find_bodies([right_gripper_name], preserve_order=True)
    
    left_gripper_pos = robot.data.body_pos_w[:, left_body_ids[0], :]
    right_gripper_pos = robot.data.body_pos_w[:, right_body_ids[0], :]
    
    gripper_center = (left_gripper_pos + right_gripper_pos) / 2.0
    distance = torch.norm(gripper_center - target_pos, dim=1) + 1e-6
    
    # ===== 添加调试打印 =====
    # 当夹爪接近手柄时（距离 < 0.15m）打印物理信息
    close_mask = distance < 0.03
    if torch.any(close_mask):
        debug_robot_physics(env, robot_cfg, "right_hand_body", distance_threshold=0.03, print_interval=5)
    # ========================
    
    position_reward = 1.0 - torch.tanh(distance / std)
    
    # 计算距离后插入
    if env.common_step_counter % 500 == 0:
        print("pos_align dbg env0: "
            f"handle={handle_pos_w[0].cpu().numpy()}, "
            f"target={target_pos[0].cpu().numpy()}, "
            f"gripL={left_gripper_pos[0].cpu().numpy()}, "
            f"gripR={right_gripper_pos[0].cpu().numpy()}, "
            f"center={gripper_center[0].cpu().numpy()}, "
            f"dist={distance[0].item():.4f}")

    return position_reward


# def gripper_handle_position_alignment(env: ManagerBasedRLEnv,
#                                     dryingbox_cfg: SceneEntityCfg = SceneEntityCfg("dryingbox"),
#                                     robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
#                                     handle_name: str = "Equipment_BB_13_C",
#                                     left_gripper_name: str = "right_hand_l",
#                                     right_gripper_name: str = "right_hand__r",
#                                     height_offset: float = 0.03,
#                                     std: float = 0.1) -> torch.Tensor:

#     dryingbox: RigidObject = env.scene[dryingbox_cfg.name]
#     robot: Articulation = env.scene[robot_cfg.name]
    
#     handle_ids, _ = dryingbox.find_bodies([handle_name], preserve_order=True)
#     handle_pos_w = dryingbox.data.body_pos_w[:, handle_ids[0], :]
#     target_pos = handle_pos_w.clone()
#     target_pos[:, 1] += height_offset
#     target_pos[:, 0] += 0.02
#     # target_pos[:, 2] -= 0.02
    
#     left_body_ids, _ = robot.find_bodies([left_gripper_name], preserve_order=True)
#     right_body_ids, _ = robot.find_bodies([right_gripper_name], preserve_order=True)
    
#     left_gripper_pos = robot.data.body_pos_w[:, left_body_ids[0], :]
#     right_gripper_pos = robot.data.body_pos_w[:, right_body_ids[0], :]
    
#     gripper_center = (left_gripper_pos + right_gripper_pos) / 2.0
#     distance = torch.norm(gripper_center - target_pos, dim=1) + 1e-6
    
#     # ===== 添加调试打印 =====
#     # 当夹爪接近手柄时（距离 < 0.15m）打印物理信息
#     close_mask = distance < 0.03
#     if torch.any(close_mask):
#         debug_robot_physics(env, robot_cfg, "right_hand_body", distance_threshold=0.03, print_interval=5)
#     # ========================
    
#     position_reward = 1.0 - torch.tanh(distance / std)
    
#     # 计算距离后插入
#     if env.common_step_counter % 00 == 0:
#         print("pos_align dbg env0: "
#             f"handle={handle_pos_w[0].cpu().numpy()}, "
#             f"target={target_pos[0].cpu().numpy()}, "
#             f"gripL={left_gripper_pos[0].cpu().numpy()}, "
#             f"gripR={right_gripper_pos[0].cpu().numpy()}, "
#             f"center={gripper_center[0].cpu().numpy()}, "
#             f"dist={distance[0].item():.4f}")

#     return position_reward


# def gripper_open_when_far(
#     env: ManagerBasedRLEnv,
#     dryingbox_cfg: SceneEntityCfg,
#     robot_cfg: SceneEntityCfg,
#     handle_name: str,
#     left_gripper_name: str,
#     right_gripper_name: str,
#     joint_names: list,
#     dist_threshold: float = 0.5,
#     target_open_pos: float = 0.3,
# ) -> torch.Tensor:
#     """
#     当夹爪距离目标点较远（> dist_threshold）时，奖励夹爪保持打开状态。
#     """
#     # 1. 获取位置
#     dryingbox = env.scene[dryingbox_cfg.name]
#     robot = env.scene[robot_cfg.name]
    
#     # 获取手柄位置
#     handle_ids, _ = dryingbox.find_bodies([handle_name])
#     handle_pos = dryingbox.data.body_pos_w[:, handle_ids[0], :3]
    
#     # 获取夹爪中心位置
#     l_ids, _ = robot.find_bodies([left_gripper_name])
#     r_ids, _ = robot.find_bodies([right_gripper_name])
#     pos_L = robot.data.body_pos_w[:, l_ids[0], :3]
#     pos_R = robot.data.body_pos_w[:, r_ids[0], :3]

#     gripper_center = (pos_L + pos_R) / 2.0
    
#     # 2. 计算距离
#     dist = torch.norm(gripper_center - handle_pos, dim=-1)
    
#     # 3. 获取手指关节位置的绝对值
#     joint_ids, _ = robot.find_joints(joint_names)
#     current_pos_abs = torch.abs(robot.data.joint_pos[:, joint_ids])
#     mean_pos_abs = torch.mean(current_pos_abs, dim=1)
    
#     # 4. 计算奖励：如果距离 > 0.5，奖励 mean_pos_abs 接近 target_open_pos (0.3)
#     # 使用 tanh 误差函数，误差越小奖励越高
#     error = torch.abs(mean_pos_abs - target_open_pos)
#     reward = 1.0 - torch.tanh(error / 0.1) # 0.1 为缩放因子，可根据需要调整
    
#     # 只有距离大于阈值时才发放奖励
#     return torch.where(dist > dist_threshold, reward, torch.zeros_like(reward))


def gripper_open_when_far(
    env: ManagerBasedRLEnv,
    dryingbox_cfg: SceneEntityCfg,   # 保留接口一致性，但不再使用
    robot_cfg: SceneEntityCfg,
    handle_name: str,               # 保留接口一致性，但不再使用
    left_gripper_name: str,
    right_gripper_name: str,
    joint_names: list,              # 不再使用
    dist_threshold: float = 0.5,    # 不再使用
    target_open_dist: float = 0.06, # 两手指完全打开时的距离（米）
) -> torch.Tensor:
    """
    奖励夹爪最大张开状态：
    只根据左右两根手指之间的实际距离计算奖励，
    当距离接近 target_open_dist (0.06m) 时奖励最大。
    """

    # 1. 获取机器人
    robot = env.scene[robot_cfg.name]

    # 2. 获取左右手指的世界坐标
    l_ids, _ = robot.find_bodies([left_gripper_name])
    r_ids, _ = robot.find_bodies([right_gripper_name])

    pos_L = robot.data.body_pos_w[:, l_ids[0], :3]
    pos_R = robot.data.body_pos_w[:, r_ids[0], :3]

    # 3. 计算两手指之间的实际距离
    finger_dist = torch.norm(pos_L - pos_R, dim=-1)

    # 4. 奖励：距离越接近最大打开距离，奖励越高
    error = torch.abs(finger_dist - target_open_dist)
    reward = 1.0 - torch.tanh(error / 0.01)  # 0.01m = 1cm 作为平滑尺度

    return reward
