# 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: {dot_fwd}")
        print(f"rew_yaw: {rew_yaw}")
        
    return total_reward


# 在 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.005
    # 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 % 500 == 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,
#     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",
#     joint_names: list = ["right_hand_joint_left", "right_hand_joint_right"],  # 保留接口但不强依赖
#     # 关键参数
#     handle_thickness: float = 0.0388,   # 把手厚度（米），≈3.88cm
#     thickness_tol: float = 0.008,       # 容差（米），≈0.8cm，控制惩罚/奖励斜率
#     grasp_radius: float = 0.04,         # 抓取判定半径，略大于把手厚度
#     height_offset: float = 0.00,        # 如需上下偏移可调整
#     target_close_pos: float = 0.03,     # 保留给 joint/clamping 兼容
# ) -> torch.Tensor:
#     dryingbox = env.scene[dryingbox_cfg.name]
#     robot = env.scene[robot_cfg.name]

#     # 目标位置（把手中心，可做微偏移）
#     handle_ids, _ = dryingbox.find_bodies([handle_name], preserve_order=True)
#     handle_pos = dryingbox.data.body_pos_w[:, handle_ids[0], :3].clone()
#     handle_pos[:, 1] += 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]
#     center = (pos_L + pos_R) / 2.0

#     # 到把手的距离（只在近距离才给分）
#     dist = torch.norm(center - handle_pos, dim=-1)
#     is_near = dist < grasp_radius

#     # 实际两指间距
#     finger_dist = torch.norm(pos_L - pos_R, dim=-1)

#     # 奖励：两指间距越接近把手厚度越好，平滑衰减
#     error = torch.abs(finger_dist - handle_thickness)
#     close_reward = 1.0 - torch.tanh(error / thickness_tol)

#     # 避免 NaN
#     close_reward = torch.nan_to_num(close_reward, nan=0.0)

#     return torch.where(is_near, close_reward, torch.zeros_like(close_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 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.005
    # 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_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


def gripper_open_penalty_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.03,   # 抓取豁免区（在此距离内不惩罚）
    # max_penalty_dist: float = 0.2,  # 这个参数现在没用了，或者可以用来控制斜率，这里简化掉
) -> torch.Tensor:
    """
    修正后的惩罚项：
    在 dist_threshold (3cm) 之外，如果夹爪没有完全张开，则给予惩罚。
    """
    robot = env.scene[robot_cfg.name]
    dryingbox = env.scene[dryingbox_cfg.name]

    # 1. 获取位置
    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
    # distance = torch.norm(gripper_center - handle_pos, dim=1)

    # 优化计算：直接用预计算的指尖中心（如果你的环境支持），或者保持原样
    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
    distance = torch.norm(gripper_center - handle_pos, dim=1)
    
    # 2. 计算闭合程度 (0=全开, 1=全闭)
    # 获取关节位置更直接，比算指尖距离更准
    joint_ids, _ = robot.find_joints(joint_names)
    joint_pos = robot.data.joint_pos[:, joint_ids]
    # 假设 0.0 是开，0.03 是闭。取平均绝对值归一化。
    # closedness 越大表示闭合越多
    closedness = torch.clamp(torch.mean(torch.abs(joint_pos), dim=1) / 0.03, 0.0, 1.0)
    
    # 3. 惩罚门控逻辑 (修正的核心)
    # 如果 distance > dist_threshold，开始惩罚
    # 使用 tanh 或 sigmoid 平滑过渡，避免硬开关导致的震荡
    # 当 dist = 0.03 时，gate = 0
    # 当 dist = 0.05 时，gate 接近 1
    is_far_gate = torch.tanh((distance - dist_threshold) * 100.0) 
    # clip 到 [0, 1]，防止 distance < threshold 时出现负数干扰
    is_far_gate = torch.clamp(is_far_gate, 0.0, 1.0)

    # 4. 计算最终惩罚
    # 只有在 "远处" 且 "闭合" 时才扣分
    penalty = is_far_gate * closedness
    
    return -penalty

# def gripper_open_penalty_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.03,   # 3cm 抓取保护区
#     max_penalty_dist: float = 0.2, # 0.55m 惩罚激活边界
# ) -> torch.Tensor:
#     """
#     使用 Sigmoid 激活函数的惩罚项：
#     1. 距离把手 < 0.55m 时，惩罚迅速激活并保持在高位 (接近 1.0)。
#     2. 距离把手 < 0.03m 时，惩罚强制清零，允许闭合抓取。
#     3. 两手指距离越小，惩罚值越大。
#     """
#     robot = env.scene[robot_cfg.name]
#     dryingbox = env.scene[dryingbox_cfg.name]

#     # 1. 获取位置信息
#     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]
    
#     # 2. 计算距离
#     gripper_center = (pos_L + pos_R) / 2.0
#     distance = torch.norm(gripper_center - handle_pos, dim=1)
#     finger_dist = torch.norm(pos_L - pos_R, dim=-1)

#     # 3. 距离激活门控 (Sigmoid 方法)
#     # 当 distance < 0.55 时，(0.55 - dist) 为正，sigmoid 趋近于 1.0
#     # 50.0 是陡峭系数，系数越大，0.55m 处的切换越像开关
#     dist_gate = torch.sigmoid((max_penalty_dist - distance) * 50.0)

#     # 4. 抓取区保护 (Deadzone)
#     # 距离小于 3cm 时，is_not_near 变为 0，完全撤销惩罚
#     is_not_near = (distance > dist_threshold).float()

#     # 5. 闭合程度计算 (0 = 全开, 1 = 全闭)
#     # 假设手指最大张开距离为 0.06m
#     max_open = 0.06
#     closedness = torch.clamp(1.0 - (finger_dist / max_open), 0.0, 1.0)

#     # 6. 综合计算并 Clamp
#     # 只有在 (0.03 < distance < 0.55) 范围内且手指闭合时，值才会很大
#     penalty_magnitude = dist_gate * is_not_near * closedness
    
#     # 返回负值作为惩罚 (值域 [-1, 0])
#     return -torch.clamp(penalty_magnitude, 0.0, 1.0)


# def gripper_open_until_near(
#     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.03,  # 3cm 阈值
# ) -> torch.Tensor:
#     """奖励夹爪在远离目标时保持打开状态，接近后不再约束。"""
#     robot = env.scene[robot_cfg.name]
#     dryingbox = env.scene[dryingbox_cfg.name]

#     # 1. 计算夹爪中心到把手的距离
#     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])
#     gripper_center = (robot.data.body_pos_w[:, l_ids[0], :3] + robot.data.body_pos_w[:, r_ids[0], :3]) / 2.0
    
#     distance = torch.norm(gripper_center - handle_pos, dim=1)
    
#     # 2. 获取夹爪当前关节位置 (假设 0 是完全打开)
#     joint_ids, _ = robot.find_joints(joint_names)
#     # 计算当前关节位置与 0（打开状态）的偏差
#     joint_pos_abs = torch.abs(robot.data.joint_pos[:, joint_ids])
#     mean_error = torch.mean(joint_pos_abs, dim=1)
    
#     # 3. 计算奖励逻辑
#     # 距离 > 3cm 时：关节越接近 0（打开），奖励越高 (最大为 1.0)
#     # 距离 <= 3cm 时：奖励为 0，不干涉模型闭合
#     is_far = distance > dist_threshold
#     open_reward = 1.0 - torch.tanh(mean_error / 0.01) # 1cm 内平滑衰减
    
#     return torch.where(is_far, open_reward, torch.zeros_like(open_reward))

# ... 在文件末尾添加以下函数 ...

def gripper_grasp_stability(
    env: ManagerBasedRLEnv,
    robot_cfg: SceneEntityCfg,
    joint_names: list,
    vel_threshold: float = 0.005  # 极低的速度阈值，判定为“静止”
) -> torch.Tensor:
    """
    奖励夹爪在闭合状态下的稳定性。
    只有当关节速度极小时（即没有在抖动或移动），才发放奖励。
    """
    robot = env.scene[robot_cfg.name]
    joint_ids, _ = robot.find_joints(joint_names)
    
    # 获取夹爪关节的绝对速度
    joint_vel = torch.abs(robot.data.joint_vel[:, joint_ids])
    
    # 判定是否所有夹爪关节都处于“静止”状态
    is_stable = torch.all(joint_vel < vel_threshold, dim=1)
    
    # 只有在闭合状态下谈稳定性才有意义，所以通常建议配合 is_grasped 使用
    # 这里我们只返回稳定性系数
    return is_stable.float()

def gripper_close_interaction_v2(
    env: ManagerBasedRLEnv,
    dryingbox_cfg: SceneEntityCfg,
    robot_cfg: SceneEntityCfg,
    handle_name: str,
    left_gripper_name: str,
    right_gripper_name: str,
    joint_names: list,
    target_close_pos: float = 0.012,
    height_offset: float = 0.00,
    grasp_radius: float = 0.05
) -> torch.Tensor:
    """
    增强版闭合奖励：结合了位置、闭合程度和物理稳定性。
    """
    # 1. 基础抓取判定
    is_grasped = _is_grasped(
        env, dryingbox_cfg, robot_cfg, handle_name,
        left_gripper_name, right_gripper_name,
        joint_names, height_offset, grasp_radius, target_close_pos
    )
    
    # 2. 计算闭合程度 (0.0 到 1.0)
    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. 计算稳定性因子 (速度越快，奖励衰减越厉害)
    joint_vel = torch.norm(robot.data.joint_vel[:, joint_ids], dim=1)
    # 使用 exp(-x) 让奖励对速度非常敏感，速度稍大奖励就迅速消失
    stability_factor = torch.exp(-20.0 * joint_vel) 
    
    # 综合奖励：必须满足抓取条件，且闭合得越深、动作越稳，分数越高
    return torch.where(is_grasped, clamping_factor * stability_factor, torch.zeros_like(clamping_factor))