# 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_lid_position_alignment(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",
#                                         height_offset: float = 0.1,
#                                         std: float = 0.1) -> torch.Tensor:
#     """奖励函数：夹爪相对于lid的位置对齐。
    
#     Args:
#         env: 环境实例
#         lid_cfg: lid的配置
#         robot_cfg: 机器人的配置
#         left_gripper_name: 左侧夹爪body名称
#         right_gripper_name: 右侧夹爪body名称
#         side_distance: Y方向的期望距离
#         height_offset: Z方向的期望高度偏移
#         std: tanh核函数的标准差
        
#     Returns:
#         位置对齐奖励 (num_envs,)
#     """
#     lid: RigidObject = env.scene[lid_cfg.name]
#     robot: Articulation = env.scene[robot_cfg.name]
    
#     lid_pos_w = lid.data.root_pos_w[:, :3]
    
#     # 获取两个夹爪的位置
#     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)
    
#     # 计算相对位置（夹爪中心相对于 lid 中心）
#     rel_pos = gripper_center - lid_pos_w  # (num_envs, 3)
    
#     # X 方向位置：应该对齐（接近 0）
#     x_dist = torch.abs(rel_pos[:, 0]) + 1e-6
#     x_reward = 1.0 - torch.tanh(x_dist / std)
    
#     # Y 方向位置：应该在 lid 的 Y 方向两侧（距离 side_distance）
#     y_dist = torch.abs(rel_pos[:, 1]) + 1e-6
#     # y_error = torch.abs(y_dist - side_distance)
#     # y_reward = 1.0 - torch.tanh(y_error / std)
#     y_reward = 1.0 - torch.tanh(y_dist / std)
    
#     # Z 方向位置：应该在 lid 上方 height_offset 处
#     # z_error = torch.clamp(torch.abs(rel_pos[:, 2] - height_offset),-3,3)+ 1e-6
#     z_error = torch.abs(rel_pos[:, 2] - height_offset) + 1e-6
#     z_reward = 1.0 - torch.tanh(z_error / std)
    
#     # 组合位置奖励
#     position_reward = x_reward * y_reward * z_reward
    
#     return position_reward


def gripper_lid_position_alignment(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",
                                    height_offset: float = 0.02,
                                    std: float = 0.1) -> torch.Tensor:
    """奖励函数：夹爪相对于lid的位置对齐（简化版）。
    
    计算夹爪中心到目标点（lid上方height_offset处）的距离奖励。
    
    Args:
        env: 环境实例
        lid_cfg: lid的配置
        robot_cfg: 机器人的配置
        left_gripper_name: 左侧夹爪body名称
        right_gripper_name: 右侧夹爪body名称
        height_offset: Z方向的期望高度偏移（目标点在lid上方）
        std: tanh核函数的标准差
        
    Returns:
        位置对齐奖励 (num_envs,)
    """
    lid: RigidObject = env.scene[lid_cfg.name]
    robot: Articulation = env.scene[robot_cfg.name]
    
    lid_pos_w = lid.data.root_pos_w[:, :3]
    
    # 目标点：lid位置 + height_offset（Z方向）
    target_pos = lid_pos_w.clone()
    # print(target_pos[0])
    target_pos[:, 2] += 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_lid_orientation_alignment(
#     env,
#     lid_cfg: SceneEntityCfg = SceneEntityCfg("lid"),
#     robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
#     gripper_body_name: str = "left_hand_body",
#     gripper_forward_axis: tuple = (0.0, 0.0, 1.0),  # 修正为 Z
#     gripper_width_axis: tuple = (0.0, 1.0, 0.0),    # 修正为 Y (用于与把手平行)
#     lid_handle_axis: tuple = (0.0, 1.0, 0.0),       # 修正为 Y
# ) -> torch.Tensor:
    
#     # 1. 获取对象
#     lid = env.scene[lid_cfg.name]
#     robot = env.scene[robot_cfg.name]
    
#     # 2. 获取姿态
#     lid_quat_w = lid.data.root_quat_w
#     body_ids, _ = robot.find_bodies([gripper_body_name], preserve_order=True)
#     gripper_quat_w = robot.data.body_quat_w[:, body_ids[0], :]

#     # 3. 垂直对齐 (Vertical Alignment)
#     # 目标：夹爪 +Y 指向 World -Z
#     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)
#     # print(grip_fwd_world[0])
#     world_down = torch.tensor([0.0, 0.0, -1.0], device=env.device).repeat(env.num_envs, 1)
    
#     dot_vertical = torch.sum(grip_fwd_world * world_down, dim=1)
#     dot_vertical = torch.clamp(dot_vertical, -1.0, 1.0)  #add
#     # 映射 [-1, 1] -> [0, 1]，且全程有梯度
#     # 即使是朝天 (+1)，只要稍微往下转一点，分数就会增加
#     # rew_vertical = (dot_vertical + 1.0) / 2.0
#     val_vertical = torch.clamp((dot_vertical + 1.0) / 2.0, 0.0, 1.0)
#     rew_vertical = torch.pow(val_vertical, 2) + 1e-4
#     # nan -> 0
#     rew_vertical = torch.where(torch.isnan(rew_vertical), torch.zeros_like(rew_vertical), rew_vertical)

#     # 4. 偏航对齐 (Yaw Alignment)
#     # 目标：夹爪 +Z 平行于 Lid +Y
#     grip_width_local = torch.tensor(gripper_width_axis, device=env.device).repeat(env.num_envs, 1)
#     grip_width_world = quat_apply(gripper_quat_w, grip_width_local)
    
#     lid_handle_local = torch.tensor(lid_handle_axis, device=env.device).repeat(env.num_envs, 1)
#     lid_handle_world = quat_apply(lid_quat_w, lid_handle_local)
    
#     dot_yaw = torch.sum(grip_width_world * lid_handle_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_vertical * 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


def gripper_lid_orientation_alignment(
    env,
    lid_cfg: SceneEntityCfg = SceneEntityCfg("lid"),
    robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
    gripper_body_name: str = "left_hand_body",
    gripper_forward_axis: tuple = (0.0, 0.0, 1.0),  # 修正为 Z
    gripper_width_axis: tuple = (0.0, 1.0, 0.0),    # 修正为 Y (用于与把手平行)
    lid_handle_axis: tuple = (0.0, 1.0, 0.0),       # 修正为 Y
    max_angle_deg: float = 60.0,  # 允许的最大角度偏差（度）
) -> torch.Tensor:
    
    # 1. 获取对象
    lid = env.scene[lid_cfg.name]
    robot = env.scene[robot_cfg.name]
    
    # 2. 获取姿态
    lid_quat_w = lid.data.root_quat_w
    body_ids, _ = robot.find_bodies([gripper_body_name], preserve_order=True)
    gripper_quat_w = robot.data.body_quat_w[:, body_ids[0], :]

    # 3. 垂直对齐 (Vertical Alignment)
    # 目标：夹爪前向轴指向 World -Z（向下）
    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)
    world_down = torch.tensor([0.0, 0.0, -1.0], device=env.device).repeat(env.num_envs, 1)
    # import pdb; pdb.set_trace()
    # 计算点积（归一化后，点积 = cos(角度)）
    dot_vertical = torch.sum(grip_fwd_world * world_down, dim=1)
    dot_vertical = torch.clamp(dot_vertical, -1.0, 1.0)
    
    # 计算角度（弧度）
    angle_rad = torch.acos(torch.clamp(dot_vertical, -1.0, 1.0))
    max_angle_rad = torch.tensor(max_angle_deg * 3.14159265359 / 180.0, device=env.device)
    
    # 如果角度 <= max_angle_deg，给予奖励
    # 奖励从角度=0时的1.0平滑衰减到角度=max_angle时的0.0
    angle_normalized = torch.clamp(angle_rad / max_angle_rad, 0.0, 1.0)
    rew_vertical = 1.0 - angle_normalized  # 角度越小，奖励越大
    
    # 如果角度超过60度，奖励为0
    rew_vertical = torch.where(angle_rad <= max_angle_rad, rew_vertical, torch.zeros_like(rew_vertical))
    
    # nan -> 0
    rew_vertical = torch.where(torch.isnan(rew_vertical), torch.zeros_like(rew_vertical), rew_vertical)

    # 4. 偏航对齐 (Yaw Alignment)
    # 目标：夹爪 +Z 平行于 Lid +Y
    grip_width_local = torch.tensor(gripper_width_axis, device=env.device).repeat(env.num_envs, 1)
    grip_width_world = quat_apply(gripper_quat_w, grip_width_local)
    
    lid_handle_local = torch.tensor(lid_handle_axis, device=env.device).repeat(env.num_envs, 1)
    lid_handle_world = quat_apply(lid_quat_w, lid_handle_local)
    
    dot_yaw = torch.sum(grip_width_world * lid_handle_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_vertical * 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


# def gripper_open_close_reward(
#     env: ManagerBasedRLEnv,
#     lid_cfg: SceneEntityCfg = SceneEntityCfg("lid"),
#     robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
#     left_gripper_name: str = "left_hand_body", # 用于计算距离的 body
#     # 注意：确保 joint_names 只包含那两个夹爪关节
#     joint_names: list = ["left_hand_joint_left", "left_hand_joint_right"], 
#     close_threshold: float = 0.05, # 5cm 以内视为“近”，需要闭合
#     target_open_pos: float = 0.0,  # 【修正】张开是 0.0
#     target_close_pos: float = 0.03 # 【修正】闭合是 0.03
# ) -> torch.Tensor:
#     """
#     逻辑很简单：
#     1. 如果 距离 < close_threshold：目标关节位置 = 0.03 (闭合)
#     2. 如果 距离 >= close_threshold：目标关节位置 = 0.0 (张开)
#     3. 返回 -abs(当前关节 - 目标关节)，即惩罚误差。
#     """
    
#     # 1. 获取对象
#     lid = env.scene[lid_cfg.name]
#     robot = env.scene[robot_cfg.name]
    
#     # 2. 计算距离 (末端 vs Lid中心)
#     lid_pos_w = lid.data.root_pos_w[:, :3]
    
#     body_ids, _ = robot.find_bodies([left_gripper_name])
#     gripper_pos_w = robot.data.body_pos_w[:, body_ids[0], :3]
    
#     # distance: (num_envs,)
#     distance = torch.norm(gripper_pos_w - lid_pos_w, dim=-1)
    
#     # 3. 动态确定目标关节值
#     # 如果距离小于阈值，目标设为 close_pos，否则设为 open_pos
#     # target_val: (num_envs, 1)
#     # print(distance)
#     target_val = torch.where(
#         distance < close_threshold,
#         torch.tensor(target_close_pos, device=env.device),
#         torch.tensor(target_open_pos, device=env.device)
#     ).unsqueeze(1)
    
#     # 4. 获取当前关节位置
#     joint_ids, _ = robot.find_joints(joint_names)
#     # joint_pos: (num_envs, 2)
#     current_joint_pos = robot.data.joint_pos[:, joint_ids]
    
#     # 取绝对值（防止左指是负数的情况，虽然你的配置里范围看起来都是正的，加上abs更保险）
#     current_joint_pos_abs = torch.abs(current_joint_pos)
    
#     # 5. 计算误差
#     # error: (num_envs,) -> 两个手指误差的平均值
#     error = torch.mean(torch.abs(current_joint_pos_abs - target_val), dim=1)
    
#     # # 6. 返回负奖励 (惩罚)
#     # # 误差越大，惩罚越大。完全符合要求时奖励为 0。
#     # return -error
#     # =====================================================
#     # 🚀 核心修复：安全输出
#     # =====================================================
#     # 1. 使用 tanh 限制数值范围在 [-1, 0]
#     reward = -torch.tanh(error / 0.01)
    
#     # 2. 过滤 NaN！这一步能救命！
#     # 如果物理引擎算出 NaN，这里会把它变成 0.0，防止训练崩溃
#     reward = torch.nan_to_num(reward, nan=0.0, posinf=0.0, neginf=-1.0)
    
#     return reward


# def gripper_open_close_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"], 
#     close_threshold: float = 0.05, 
#     target_open_pos: float = 0.0,
#     target_close_pos: float = 0.03,
#     height_offset: float = 0.06,  # 新增：与位置奖励保持一致
# ) -> torch.Tensor:
#     """
#     逻辑：
#     1. 计算夹爪到【目标抓取点】(Lid上方 height_offset 处) 的距离。
#     2. 如果距离 < close_threshold：目标关节位置 = 闭合。
#     3. 否则：目标关节位置 = 张开。
#     """
    
#     # 1. 获取对象
#     lid = env.scene[lid_cfg.name]
#     robot = env.scene[robot_cfg.name]
    
#     # 2. 计算目标抓取点位置 (Lid位置 + Z轴偏移)
#     lid_pos_w = lid.data.root_pos_w[:, :3].clone()
#     lid_pos_w[:, 2] += height_offset # 修正为抓取点位置
    
#     # 获取夹爪位置
#     body_ids, _ = robot.find_bodies([left_gripper_name])
#     gripper_pos_w = robot.data.body_pos_w[:, body_ids[0], :3]
    
#     # 3. 计算距离 (夹爪 vs 目标抓取点)
#     distance = torch.norm(gripper_pos_w - lid_pos_w, dim=-1)
    
#     # 4. 动态确定目标关节值
#     # 靠近抓取点 -> 闭合；远离 -> 张开
#     target_val = torch.where(
#         distance < close_threshold,
#         torch.tensor(target_close_pos, device=env.device),
#         torch.tensor(target_open_pos, device=env.device)
#     ).unsqueeze(1)
    
#     # 5. 获取当前关节位置
#     joint_ids, _ = robot.find_joints(joint_names)
#     current_joint_pos = robot.data.joint_pos[:, joint_ids]
#     current_joint_pos_abs = torch.abs(current_joint_pos)
    
#     # 6. 计算误差并返回奖励
#     error = torch.mean(torch.abs(current_joint_pos_abs - target_val), dim=1)
    
#     # 使用 tanh 限制数值范围
#     reward = 1.0 - torch.tanh(error / 0.01) # 误差越小奖励越大
    
#     return torch.nan_to_num(reward, nan=0.0)


def _is_grasped(
    env: ManagerBasedRLEnv,
    lid_cfg: SceneEntityCfg,
    robot_cfg: SceneEntityCfg,
    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. (关键) 夹爪Z轴高度合适，不能压在盖子上面 (Z-Axis Check)
    """
    # 1. 获取对象
    lid = env.scene[lid_cfg.name]
    robot = env.scene[robot_cfg.name]
    
    # 2. 目标点位置 (把手中心)
    lid_pos_w = lid.data.root_pos_w[:, :3].clone()
    lid_pos_w[:, 2] += height_offset 
    
    # 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 - lid_pos_w, dim=-1)
    is_near = dist_to_handle < grasp_radius
    
    # 5. 条件二：Z轴高度判定 (防止压在盖子上)
    # 夹爪中心必须在把手目标点附近，允许上下微小浮动，但不能太高
    # 假设 height_offset 是把手几何中心，那么夹爪不应该比它高太多
    z_diff = gripper_center[:, 2] - lid_pos_w[:, 2]
    is_z_valid = torch.abs(z_diff) < 0.03 # 3cm 容差
    
    # 6. 条件三：闭合判定 (正在尝试闭合)
    joint_ids, _ = robot.find_joints(joint_names)
    # 获取关节位置 (绝对值)
    joint_pos = torch.abs(robot.data.joint_pos[:, joint_ids])
    
    # === 修改闭合判定 ===
    
    # 1. 关节位置判定 (Effort Check)
    # 只要用力了就行，去掉上限判定，防止夹太紧被误判
    # 假设 0.05 是你的新 target_close_pos
    is_effort_closing = torch.all(joint_pos > 0.005, dim=1)
    
    # 2. 几何排斥判定 (Geometry Check)
    # 如果真的夹住了东西，左右指尖的距离应该被把手撑开，不会太小
    # 获取指尖位置
    dist_fingertips = torch.norm(pos_L - pos_R, dim=-1)
    
    # 假设把手厚度是 1cm (0.01m)
    # 如果指尖距离 < 0.005，说明两个手指已经碰到一起了（空夹），没夹住东西
    # 如果指尖距离 >= 0.005，说明中间有东西挡着
    is_not_empty = dist_fingertips > 0.005
    
    # 综合判定：
    # 1. 在把手附近 (is_near)
    # 2. 高度合适 (is_z_valid)
    # 3. 在用力闭合 (is_effort_closing)
    # 4. 指尖没贴死 (is_not_empty) -> 这就证明夹住东西了！
    return is_near & is_z_valid & is_effort_closing & is_not_empty


def gripper_close_when_near(
    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"], 
    target_close_pos: float = 0.03,
    height_offset: float = 0.02,
    grasp_radius: 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. 计算夹紧程度 (Clamping Intensity)
    robot = env.scene[robot_cfg.name]
    joint_ids, _ = robot.find_joints(joint_names)
    # 获取当前关节位置绝对值 (num_envs, num_joints)
    current_joint_pos = torch.abs(robot.data.joint_pos[:, joint_ids])
    # 计算两个指关节的平均闭合深度
    mean_joint_pos = torch.mean(current_joint_pos, dim=1)
    
    # 计算奖励系数：当前位置 / 目标闭合位置
    # 这样当关节越接近 0.03 时，奖励越接近 1.0
    # 强制让 Agent 产生“压满”动作的冲动
    clamping_factor = torch.clamp(mean_joint_pos / target_close_pos, max=1.0)
    
    # 3. 只有在满足抓取判定时，才发放带有权重的夹紧奖励
    # 如果没对准或者没夹到，奖励依然是 0
    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)