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

"""MindRobot left arm IK-Rel environment config for teleoperation."""

from __future__ import annotations
from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR
import torch

import isaaclab.sim as sim_utils
from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg
from isaaclab.controllers.differential_ik_cfg import DifferentialIKControllerCfg
from isaaclab.devices import DevicesCfg
from isaaclab.devices.keyboard import Se3KeyboardCfg
from isaaclab.devices.spacemouse import Se3SpaceMouseCfg
from isaaclab.devices.gamepad import Se3GamepadCfg
from isaaclab.envs import ManagerBasedRLEnvCfg
from isaaclab.envs.mdp.actions.actions_cfg import (
    BinaryJointPositionActionCfg,
    DifferentialInverseKinematicsActionCfg,
    JointVelocityActionCfg,
)
from isaaclab.managers import EventTermCfg as EventTerm
from isaaclab.managers import ObservationGroupCfg as ObsGroup
from isaaclab.managers import ObservationTermCfg as ObsTerm
from isaaclab.managers import TerminationTermCfg as DoneTerm
from isaaclab.scene import InteractiveSceneCfg
from isaaclab.sensors.frame_transformer.frame_transformer_cfg import FrameTransformerCfg
from isaaclab.sensors.frame_transformer.frame_transformer_cfg import OffsetCfg
from isaaclab.sim.schemas.schemas_cfg import RigidBodyPropertiesCfg
from isaaclab.sim.spawners.from_files.from_files_cfg import GroundPlaneCfg
from isaaclab.utils import configclass
from isaaclab.markers.config import FRAME_MARKER_CFG
from isaaclab.devices.openxr import XrCfg

# 导入基础配置和MDP函数
from isaaclab_tasks.manager_based.manipulation.stack import mdp

# 导入 MindRobot 资产
from .mindrobot_cfg import (
    MINDBOT_HIGH_PD_CFG,
    MINDBOT_LEFT_ARM_JOINTS,
    MINDBOT_LEFT_EEF_BODY,
    MINDBOT_LEFT_GRIPPER_JOINTS,
    MINDBOT_WHEEL_JOINTS,
    MINDBOT_LEFT_ARM_PRIM_ROOT,
    MINDBOT_LEFT_EEF_PRIM,
    MINDBOT_RIGHT_ARM_PRIM_ROOT,
)

# 在文件开头添加
import isaaclab.utils.assets as assets_utils
from mindbot.utils.assets import MINDBOT_ASSETS_DIR
from mindbot.assets.lab import ROOM_CFG
from mindbot.assets.table import TABLE_CFG
from mindbot.assets.dryingbox import DRYINGBOX_CFG
# # 然后在 scene 配置中使用
# spawn=sim_utils.UsdFileCfg(
#     usd_path=f"{assets_utils.ISAAC_NUCLEUS_DIR}/Props/Mounts/SeattleLabTable/table_instanceable.usd"
# ),
# =====================================================================
# Scene Configuration
# =====================================================================


@configclass
class MindRobotTeleopSceneCfg(InteractiveSceneCfg):
    """Minimal scene for MindRobot teleoperation: robot + table + optional cube."""

    # Ground plane
    plane = AssetBaseCfg(
        prim_path="/World/GroundPlane",
        init_state=AssetBaseCfg.InitialStateCfg(pos=[0.0, 0.0, 0.0]),
        spawn=GroundPlaneCfg(),
    )

    # # Table
    # table = AssetBaseCfg(
    #     prim_path="{ENV_REGEX_NS}/Table",
    #     init_state=AssetBaseCfg.InitialStateCfg(pos=[0.5, 0, 0], rot=[0.707, 0, 0, 0.707]),
    #     spawn=sim_utils.UsdFileCfg(
    #         usd_path=f"{assets_utils.ISAAC_NUCLEUS_DIR}/Props/Mounts/SeattleLabTable/table_instanceable.usd"
    #     ),
    # )

    # Optional: Single cube for testing (can be removed if not needed)
    # cube = RigidObjectCfg(
    #     prim_path="{ENV_REGEX_NS}/Cube",
    #     spawn=sim_utils.CuboidCfg(
    #         size=(0.04, 0.04, 0.04),
    #         rigid_props=RigidBodyPropertiesCfg(),
    #         mass_props=sim_utils.MassPropertiesCfg(mass=0.1),
    #         collision_props=sim_utils.CollisionPropertiesCfg(),
    #         visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0)),
    #     ),
    #     init_state=RigidObjectCfg.InitialStateCfg(pos=(0.5, 0.0, 0.45)),
    # )
    
    room = ROOM_CFG.replace(prim_path="{ENV_REGEX_NS}/Room")
    table = TABLE_CFG.replace(prim_path="{ENV_REGEX_NS}/Table")
    drying_box = DRYINGBOX_CFG.replace(prim_path="{ENV_REGEX_NS}/DryingBox")
    # room = AssetBaseCfg(
    #     prim_path="{ENV_REGEX_NS}/Room",
    #     spawn=sim_utils.UsdFileCfg(
    #         usd_path=f"{MINDBOT_ASSETS_DIR}/twinlab/Collected_lab2/lab.usd",
    #     ),
    # )

    # MindRobot
    robot: ArticulationCfg = MINDBOT_HIGH_PD_CFG.replace(
        prim_path="{ENV_REGEX_NS}/Robot"
    )

    # Lighting
    light = AssetBaseCfg(
        prim_path="/World/light",
        spawn=sim_utils.DomeLightCfg(color=(0.75, 0.75, 0.75), intensity=3000.0),
    )

    # End-effector frame transformer
    ee_frame: FrameTransformerCfg = None  # Will be set in __post_init__


# =====================================================================
# Actions Configuration
# =====================================================================


@configclass
class MindRobotTeleopActionsCfg:
    """Actions for MindRobot left arm IK teleoperation."""

    # Left arm IK control
    arm_action = DifferentialInverseKinematicsActionCfg(
        asset_name="robot",
        joint_names=[MINDBOT_LEFT_ARM_JOINTS],
        body_name=MINDBOT_LEFT_EEF_BODY,
        controller=DifferentialIKControllerCfg(
            command_type="pose",
            use_relative_mode=True,
            ik_method="dls",
        ),
        scale=1,
        body_offset=DifferentialInverseKinematicsActionCfg.OffsetCfg(
            pos=[0.0, 0.0, 0.0],
        ),
    )

    # Wheel velocity control for differential drive (skid-steer).
    # Joint order in articulation: right_b, left_b, left_f, right_f
    # Action vector: [right_b_vel, left_b_vel, left_f_vel, right_f_vel] in rad/s.
    wheel_action = JointVelocityActionCfg(
        asset_name="robot",
        joint_names=MINDBOT_WHEEL_JOINTS,
        scale=1.0,
    )

    # Left gripper control (binary: open/close)
    gripper_action = BinaryJointPositionActionCfg(
        asset_name="robot",
        joint_names=MINDBOT_LEFT_GRIPPER_JOINTS,
        open_command_expr={
            "left_hand_joint_left": 0.0,
            "left_hand_joint_right": 0.0,
        },
        close_command_expr={
            "left_hand_joint_left": -0.03,
            "left_hand_joint_right": 0.03,
        },
    )


# =====================================================================
# Observations Configuration
# =====================================================================


@configclass
class MindRobotTeleopObsCfg:
    """Minimal observations for MindRobot teleoperation."""

    @configclass
    class PolicyCfg(ObsGroup):
        """Observations for policy group - only robot state, no cube dependencies."""

        actions = ObsTerm(func=mdp.last_action)
        joint_pos = ObsTerm(func=mdp.joint_pos_rel)
        joint_vel = ObsTerm(func=mdp.joint_vel_rel)
        eef_pos = ObsTerm(func=mdp.ee_frame_pos)
        eef_quat = ObsTerm(func=mdp.ee_frame_quat)
        gripper_pos = ObsTerm(func=mdp.gripper_pos)

        def __post_init__(self):
            self.enable_corruption = False
            self.concatenate_terms = False

    # observation groups
    policy: PolicyCfg = PolicyCfg()


# =====================================================================
# Terminations Configuration
# =====================================================================


@configclass
class MindRobotTeleopTerminationsCfg:
    """Minimal terminations for teleoperation - only time_out or none."""

    # Optional: Keep time_out for safety, or remove entirely
    time_out = DoneTerm(func=mdp.time_out, time_out=True)


# =====================================================================
# Events Configuration
# =====================================================================


def _disable_arm_gravity(env, env_ids: torch.Tensor):
    """Disable gravity for both arm subtrees; chassis/wheels/trunk keep gravity.

    Called once at startup. The @apply_nested decorator on
    modify_rigid_body_properties recurses into all rigid-body children under
    the given prim path, so every link of the arm (including gripper) is covered.

    Arm prim roots (per env):
        Robot/rm_65_fb_left   — left RM-65 arm + gripper
        Robot/rm_65_fb_right  — right RM-65 arm + gripper
    """
    import isaaclab.sim.schemas as schemas

    arm_cfg = RigidBodyPropertiesCfg(disable_gravity=True)
    for env_id in range(env.num_envs):
        for arm_path in [
            f"/World/envs/env_{env_id}/Robot/{MINDBOT_LEFT_ARM_PRIM_ROOT}",
            f"/World/envs/env_{env_id}/Robot/{MINDBOT_RIGHT_ARM_PRIM_ROOT}",
        ]:
            schemas.modify_rigid_body_properties(arm_path, arm_cfg)


@configclass
class MindRobotTeleopEventsCfg:
    """Reset events for teleoperation: R键重置时将场景恢复到初始状态。"""

    disable_arm_gravity = EventTerm(
        func=_disable_arm_gravity,
        mode="startup",
    )

    reset_scene = EventTerm(
        func=mdp.reset_scene_to_default,
        mode="reset",
    )


# =====================================================================
# Main Environment Configuration
# =====================================================================


@configclass
class EmptyCfg:
    pass


@configclass
class MindRobotLeftArmIKEnvCfg(ManagerBasedRLEnvCfg):
    """MindRobot 左臂 IK-Rel 遥操作环境配置（最小化版本）。"""

    # Scene settings
    scene: MindRobotTeleopSceneCfg = MindRobotTeleopSceneCfg(
        num_envs=1,
        env_spacing=2.5,
    )

    # Basic settings
    observations: MindRobotTeleopObsCfg = MindRobotTeleopObsCfg()
    actions: MindRobotTeleopActionsCfg = MindRobotTeleopActionsCfg()

    # MDP settings
    terminations: MindRobotTeleopTerminationsCfg = MindRobotTeleopTerminationsCfg()
    events: MindRobotTeleopEventsCfg = MindRobotTeleopEventsCfg()

    # Unused managers
    commands: EmptyCfg = EmptyCfg()
    rewards: EmptyCfg = EmptyCfg()
    curriculum: EmptyCfg = EmptyCfg()

    # XR configuration for hand tracking (if needed)
    xr: XrCfg = XrCfg(
        anchor_pos=(-0.1, -0.5, -1.05),
        anchor_rot=(0.866, 0, 0, -0.5),
    )

    def __post_init__(self):
        """Post initialization."""
        super().__post_init__()

        # General settings
        self.decimation = 2
        self.episode_length_s = 50.0

        # Simulation settings
        self.sim.dt = 0.01  # 100Hz
        self.sim.render_interval = 2

        # Keep the base mobile for teleoperation.
        # The absolute arm IK command must therefore be expressed in the robot
        # root frame instead of assuming a fixed world-aligned base.
        robot_cfg = MINDBOT_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
        robot_cfg.spawn.articulation_props.fix_root_link = False
        self.scene.robot = robot_cfg

        # Configure end-effector frame transformer
        marker_cfg = FRAME_MARKER_CFG.copy()
        marker_cfg.markers["frame"].scale = (0.1, 0.1, 0.1)
        marker_cfg.prim_path = "/Visuals/FrameTransformer"
        self.scene.ee_frame = FrameTransformerCfg(
            prim_path=f"{{ENV_REGEX_NS}}/Robot/{MINDBOT_LEFT_ARM_PRIM_ROOT}/base_link",
            debug_vis=False,
            visualizer_cfg=marker_cfg,
            target_frames=[
                FrameTransformerCfg.FrameCfg(
                    prim_path=f"{{ENV_REGEX_NS}}/Robot/{MINDBOT_LEFT_EEF_PRIM}",
                    name="end_effector",
                    offset=OffsetCfg(pos=[0.0, 0.0, 0.0]),
                ),
            ],
        )


        # Configure teleoperation devices
        self.teleop_devices = DevicesCfg(
            devices={
                "keyboard": Se3KeyboardCfg(
                    pos_sensitivity=0.05,
                    rot_sensitivity=0.05,
                    sim_device=self.sim.device,
                ),
                "spacemouse": Se3SpaceMouseCfg(
                    pos_sensitivity=0.05,
                    rot_sensitivity=0.05,
                    sim_device=self.sim.device,
                ),
                "gamepad": Se3GamepadCfg(
                    pos_sensitivity=0.1,
                    rot_sensitivity=0.1,
                    sim_device=self.sim.device,
                ),
            }
        )

        # Gripper parameters for status checking
        self.gripper_joint_names = MINDBOT_LEFT_GRIPPER_JOINTS
        self.gripper_open_val = 0.0
        self.gripper_threshold = 0.005


@configclass
class MindRobotLeftArmIKAbsEnvCfg(MindRobotLeftArmIKEnvCfg):
    """MindRobot 左臂 IK-Abs 遥操作环境配置（绝对坐标系，适合 VR 控制器）。"""

    def __post_init__(self):
        super().__post_init__()
        # Switch controller to absolute mode
        self.actions.arm_action.controller.use_relative_mode = False
