# Copyright (c) 2024, MindRobot Project
# SPDX-License-Identifier: BSD-3-Clause

"""Custom observation functions for open_drybox teleoperation environments."""

from __future__ import annotations

from typing import TYPE_CHECKING

import torch

from isaaclab.managers import SceneEntityCfg

if TYPE_CHECKING:
    from isaaclab.envs import ManagerBasedRLEnv


def gripper_pos(
    env: ManagerBasedRLEnv,
    robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
    joint_names: list[str] = None,
) -> torch.Tensor:
    """Return the opening state of a parallel gripper specified by joint names.

    Unlike the built-in ``gripper_pos`` which reads from ``env.cfg.gripper_joint_names``,
    this version accepts ``joint_names`` directly so the same function can be used
    for both left and right grippers in a dual-arm setup.

    Args:
        robot_cfg: Scene entity config for the robot articulation.
        joint_names: List of exactly 2 joint names ``[joint_open_positive, joint_open_negative]``.
            The first joint increases with opening, the second decreases.

    Returns:
        Tensor of shape ``(num_envs, 2)`` with the joint positions (sign-corrected).
    """
    if joint_names is None:
        raise ValueError("joint_names must be specified")

    robot = env.scene[robot_cfg.name]
    ids, _ = robot.find_joints(joint_names)

    if len(ids) != 2:
        raise ValueError(f"Expected exactly 2 gripper joints, got {len(ids)} for {joint_names}")

    j1 = robot.data.joint_pos[:, ids[0]].clone().unsqueeze(1)
    j2 = -1.0 * robot.data.joint_pos[:, ids[1]].clone().unsqueeze(1)
    return torch.cat((j1, j2), dim=1)
