from __future__ import annotations
import torch
from typing import TYPE_CHECKING
from isaaclab.assets import Articulation, RigidObject
from isaaclab.managers import SceneEntityCfg

if TYPE_CHECKING:
    from isaaclab.envs import ManagerBasedRLEnv

def lid_dropped(env: ManagerBasedRLEnv,
                minimum_height: float = -0.05,
                lid_cfg: SceneEntityCfg = SceneEntityCfg("DexCube")) -> torch.Tensor:
    lid: RigidObject = env.scene[lid_cfg.name]
    return lid.data.root_pos_w[:, 2] < minimum_height

def lid_successfully_grasped(env: ManagerBasedRLEnv,
                             distance_threshold: float = 0.03,
                             height_threshold: float = 0.2,
                             lid_cfg: SceneEntityCfg = SceneEntityCfg("DexCube"),
                             robot_cfg: SceneEntityCfg = SceneEntityCfg("Mindbot"),
                             gripper_body_name: str = "left_hand_body") -> torch.Tensor:
    lid: RigidObject = env.scene[lid_cfg.name]
    robot: Articulation = env.scene[robot_cfg.name]
    body_idx = robot.find_bodies(gripper_body_name)[0]
    gripper_pos_w = robot.data.body_pos_w[:, body_idx, :]
    distance = torch.norm(lid.data.root_pos_w[:, :3] - gripper_pos_w, dim=1)
    is_close = distance < distance_threshold
    is_lifted = lid.data.root_pos_w[:, 2] > height_threshold
    return is_close & is_lifted

def gripper_at_lid_side(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",
                        side_distance: float = 0.05,
                        side_tolerance: float = 0.01,
                        alignment_tolerance: float = 0.02,
                        height_offset: float = 0.1,
                        height_tolerance: float = 0.02) -> torch.Tensor:
    """Terminate when gripper center is positioned on the side of the lid at specified height.
    
    坐标系说明：
    - X 方向：两个夹爪朝中心合并的方向
    - Y 方向：夹爪间空隙的方向，和 lid 的把手方向一致
    - Z 方向：高度方向
    """
    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], :]
    right_gripper_pos = robot.data.body_pos_w[:, right_body_ids[0], :]
    
    # 计算夹爪中心位置
    gripper_center = (left_gripper_pos + right_gripper_pos) / 2.0
    rel_pos = gripper_center - lid_pos_w
    
    # Y 方向：应该在 lid 的 Y 方向两侧（距离 side_distance）
    y_dist = torch.abs(rel_pos[:, 1])
    y_ok = (y_dist >= side_distance - side_tolerance) & (y_dist <= side_distance + side_tolerance)
    
    # X 方向：应该对齐（接近 0）
    x_dist = torch.abs(rel_pos[:, 0])
    x_ok = x_dist < alignment_tolerance
    
    # Z 方向：应该在 lid 上方 height_offset 处
    z_error = torch.abs(rel_pos[:, 2] - height_offset)
    z_ok = z_error < height_tolerance
    
    # 所有条件都要满足
    return x_ok & y_ok & z_ok


def base_height_failure(env: ManagerBasedRLEnv,
                        minimum_height: float | None = None,
                        maximum_height: float | None = None,
                        asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
    """Terminate when the robot's base height is outside the specified range."""
    # extract the used quantities (to enable type-hinting)
    asset: Articulation = env.scene[asset_cfg.name]
    root_pos_z = asset.data.root_pos_w[:, 2]
    
    # check if height is outside the range
    out_of_bounds = torch.zeros_like(root_pos_z, dtype=torch.bool)
    if minimum_height is not None:
        out_of_bounds |= root_pos_z < minimum_height
    if maximum_height is not None:
        out_of_bounds |= root_pos_z > maximum_height
        
    return out_of_bounds

