# scripts/environments/teleoperation/mindrobot_keyboard.py
# Copyright (c) 2022-2026, The Isaac Lab Project Developers.
# SPDX-License-Identifier: BSD-3-Clause
"""Shared keyboard controllers for MindRobot teleoperation and demo recording.

This module contains ONLY class definitions (no argparse, no AppLauncher),
so it is safe to import from both standalone scripts and from record_demos.py.
"""

import weakref
import numpy as np
import torch
import carb

from isaaclab.devices import Se3Keyboard, Se3KeyboardCfg


class WheelKeyboard:
    """Differential-drive (skid-steer) keyboard controller.

    Listens to arrow keys via Carb and produces a 4-D joint-velocity
    command: [right_b, left_b, left_f, right_f]  (rad/s).

    Key mappings
    ─────────────────────────────────────────────────────
    ↑  (UP)    forward   [ v,  v,  v,  v]
    ↓  (DOWN)  backward  [-v, -v, -v, -v]
    ←  (LEFT)  left turn [ v, -v, -v,  v]
    →  (RIGHT) right turn[-v,  v,  v, -v]
    ─────────────────────────────────────────────────────
    """

    _WHEEL_KEYS = {"UP", "DOWN", "LEFT", "RIGHT"}

    def __init__(self, wheel_speed: float = 5.0, sim_device: str = "cuda:0"):
        self.wheel_speed = wheel_speed
        self._sim_device = sim_device
        self._wheel_vel = np.zeros(4)
        self._key_map: dict[str, np.ndarray] = {}
        self._create_bindings()

        import omni
        self._appwindow = omni.appwindow.get_default_app_window()
        self._input = carb.input.acquire_input_interface()
        self._keyboard = self._appwindow.get_keyboard()
        self._keyboard_sub = self._input.subscribe_to_keyboard_events(
            self._keyboard,
            lambda event, *args, obj=weakref.proxy(self): obj._on_keyboard_event(event, *args),
        )

    def __del__(self):
        self._input.unsubscribe_to_keyboard_events(self._keyboard, self._keyboard_sub)

    def __str__(self) -> str:
        return (
            "WheelKeyboard (skid-steer chassis controller)\n"
            "\t↑ UP    — forward\n"
            "\t↓ DOWN  — backward\n"
            "\t← LEFT  — left turn (in-place)\n"
            "\t→ RIGHT — right turn (in-place)"
        )

    def reset(self):
        self._wheel_vel = np.zeros(4)

    def advance(self) -> torch.Tensor:
        return torch.tensor(self._wheel_vel.copy(), dtype=torch.float32, device=self._sim_device)

    def _create_bindings(self):
        v = self.wheel_speed
        self._key_map = {
            "UP":    np.array([ v,  v,  v,  v]),
            "DOWN":  np.array([-v, -v, -v, -v]),
            "LEFT":  np.array([ v, -v, -v,  v]),
            "RIGHT": np.array([-v,  v,  v, -v]),
        }

    def _on_keyboard_event(self, event, *args, **kwargs):
        if event.type == carb.input.KeyboardEventType.KEY_PRESS:
            key = event.input.name
            if key in self._key_map:
                self._wheel_vel += self._key_map[key]
        if event.type == carb.input.KeyboardEventType.KEY_RELEASE:
            key = event.input.name
            if key in self._key_map:
                self._wheel_vel -= self._key_map[key]
        return True


class MindRobotCombinedKeyboard:
    """组合遥操作控制器：左臂 IK (Se3Keyboard 7D) + 底盘轮速 (WheelKeyboard 4D)

    输出 11 维 action: [arm_IK(6) | wheel_vel(4) | gripper(1)]
    与 MindRobotTeleopActionsCfg 中声明顺序一致。

    Key bindings
    ─────────────────────────────────────────────
    W/S/A/D/Q/E   — arm IK translate X/Y/Z
    Z/X T/G C/V   — arm IK rotate roll/pitch/yaw
    K             — gripper toggle open/close
    ↑↓←→          — chassis forward/backward/turn
    R             — reset
    ─────────────────────────────────────────────
    """

    def __init__(
        self,
        pos_sensitivity: float = 0.05,
        rot_sensitivity: float = 0.05,
        wheel_speed: float = 5.0,
        sim_device: str = "cuda:0",
    ):
        self._arm = Se3Keyboard(
            Se3KeyboardCfg(
                pos_sensitivity=pos_sensitivity,
                rot_sensitivity=rot_sensitivity,
                sim_device=sim_device,
            )
        )
        self._wheel = WheelKeyboard(wheel_speed=wheel_speed, sim_device=sim_device)

    def reset(self):
        self._arm.reset()
        self._wheel.reset()

    def add_callback(self, key: str, callback):
        """Delegate callbacks to the arm keyboard (Se3Keyboard supports arbitrary key callbacks)."""
        self._arm.add_callback(key, callback)

    def advance(self) -> torch.Tensor:
        arm_cmd  = self._arm.advance()       # (7,): [dx,dy,dz,rx,ry,rz, gripper]
        arm_ik   = arm_cmd[:6]               # (6,)
        gripper  = arm_cmd[6:7]              # (1,)
        wheel_vel = self._wheel.advance()    # (4,)
        return torch.cat([arm_ik, wheel_vel, gripper])  # (11,)

    def __str__(self) -> str:
        return (
            "MindRobotCombinedKeyboard\n"
            "  Arm (Se3Keyboard): W/S/A/D/Q/E + Z/X/T/G/C/V, K=gripper\n"
            "  Wheel (arrows):    ↑↓←→\n"
            "  R = reset  |  L = save-success-and-reset"
        )