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

"""MindRobot dual-arm IK absolute environment config for XR teleoperation."""

from __future__ import annotations

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.devices.openxr import XrCfg
from isaaclab.envs import ManagerBasedRLEnvCfg
from isaaclab.envs.mdp.actions.actions_cfg import (
    BinaryJointPositionActionCfg,
    DifferentialInverseKinematicsActionCfg,
    JointPositionActionCfg,
    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 SceneEntityCfg
from isaaclab.managers import TerminationTermCfg as DoneTerm
from isaaclab.markers.config import FRAME_MARKER_CFG
from isaaclab.scene import InteractiveSceneCfg
from isaaclab.sensors.frame_transformer.frame_transformer_cfg import FrameTransformerCfg, OffsetCfg
from isaaclab.sensors import CameraCfg
from isaaclab.sim.schemas.schemas_cfg import RigidBodyPropertiesCfg
from isaaclab.sim.spawners.from_files.from_files_cfg import GroundPlaneCfg
from isaaclab.sim.spawners.materials import RigidBodyMaterialCfg
from isaaclab.utils import configclass
from isaaclab.envs import ViewerCfg

from . import mdp

from .mindrobot_2i_cfg import (
    MINDBOT_HIGH_PD_CFG,
    MINDBOT_LEFT_ARM_JOINTS,
    MINDBOT_RIGHT_ARM_JOINTS,
    MINDBOT_LEFT_EEF_BODY,
    MINDBOT_RIGHT_EEF_BODY,
    MINDBOT_LEFT_GRIPPER_JOINTS,
    MINDBOT_RIGHT_GRIPPER_JOINTS,
    MINDBOT_WHEEL_JOINTS,
    MINDBOT_HEAD_JOIONTS,
    MINDBOT_TRUNK_JOINT,
    MINDBOT_LEFT_ARM_PRIM_ROOT,
    MINDBOT_RIGHT_ARM_PRIM_ROOT,
    MINDBOT_LEFT_EEF_PRIM,
    MINDBOT_RIGHT_EEF_PRIM,
)

from mindbot.assets.lab import ROOM_CFG
from mindbot.assets.table import TABLE_CFG
from mindbot.assets.dryingbox import DRYINGBOX_CFG

# =====================================================================
# Scene Configuration
# =====================================================================

@configclass
class MindRobotDualArmSceneCfg(InteractiveSceneCfg):
    """Scene for MindRobot dual-arm teleoperation."""

    plane = AssetBaseCfg(
        prim_path="/World/GroundPlane",
        init_state=AssetBaseCfg.InitialStateCfg(pos=[0.0, 0.0, 0.0]),
        spawn=GroundPlaneCfg(
            # Moderate friction: enough traction to move forward, low enough to
            # allow skid-steer lateral wheel slip for turning.
            # Default (1.0/1.0) moves but can't turn; 0.4/0.3 spins in place.
            physics_material=RigidBodyMaterialCfg(
                static_friction=0.7,
                dynamic_friction=0.5,
                restitution=0.0,
            )
        ),
    )
    
    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")

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

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

    # EEF frame transformers — set in __post_init__
    ee_frame: FrameTransformerCfg = None
    ee_frame_right: FrameTransformerCfg = None

    # VR stereo cameras — reuse existing ZED_X CameraLeft/CameraRight prims from USD
    # (spawn=None means use existing prim, don't create a new one)
    # The ZED_X baseline in USD is 120mm; for perfect VR IPD we'd ideally set 65mm,
    # but reusing existing prims avoids spawn issues. stereoOffset=0 in video_source.yml
    # since sim rendering gives correct perspective regardless of physical baseline.
    vr_left_eye: CameraCfg = CameraCfg(
        prim_path="{ENV_REGEX_NS}/Robot/robot_head/ZED_X/base_link/ZED_X/CameraLeft",
        update_period=1 / 30,
        height=1440,  # 2560x1440 per eye → SBS 5120x1440; effective ~1920x1440 after shader crop
        width=2560,
        data_types=["rgb"],
        spawn=None,
    )
    vr_right_eye: CameraCfg = CameraCfg(
        prim_path="{ENV_REGEX_NS}/Robot/robot_head/ZED_X/base_link/ZED_X/CameraRight",
        update_period=1 / 30,
        height=1440,
        width=2560,
        data_types=["rgb"],
        spawn=None,
    )


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

@configclass
class MindRobotDualArmActionsCfg:
    """Actions for MindRobot dual-arm IK teleoperation (absolute mode).

    Action vector layout (total 22D):
        left_arm_action     (7D): [x, y, z, qw, qx, qy, qz]  absolute EEF pose, root frame
        wheel_action        (4D): [right_b, left_b, left_f, right_f] wheel vel (rad/s)
        left_gripper_action (1D): +1=open, -1=close
        right_arm_action    (7D): [x, y, z, qw, qx, qy, qz]  absolute EEF pose, root frame
        right_gripper_action(1D): +1=open, -1=close
        head_action         (2D): [yaw_rad, pitch_rad]  absolute joint position
    """

    left_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=False,
            ik_method="dls",
        ),
        scale=1,
        body_offset=DifferentialInverseKinematicsActionCfg.OffsetCfg(pos=[0.0, 0.0, 0.0]),
    )

    wheel_action = JointVelocityActionCfg(
        asset_name="robot",
        joint_names=MINDBOT_WHEEL_JOINTS,
        scale=1.0,
    )

    left_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},
    )

    right_arm_action = DifferentialInverseKinematicsActionCfg(
        asset_name="robot",
        joint_names=[MINDBOT_RIGHT_ARM_JOINTS],
        body_name=MINDBOT_RIGHT_EEF_BODY,
        controller=DifferentialIKControllerCfg(
            command_type="pose",
            use_relative_mode=False,
            ik_method="dls",
        ),
        scale=1,
        body_offset=DifferentialInverseKinematicsActionCfg.OffsetCfg(pos=[0.0, 0.0, 0.0]),
    )

    right_gripper_action = BinaryJointPositionActionCfg(
        asset_name="robot",
        joint_names=MINDBOT_RIGHT_GRIPPER_JOINTS,
        open_command_expr={"right_hand_joint_left": 0.0, "right_hand_joint_right": 0.0},
        close_command_expr={"right_hand_joint_left": -0.03, "right_hand_joint_right": 0.03},
    )

    head_action = JointPositionActionCfg(
        asset_name="robot",
        joint_names=MINDBOT_HEAD_JOIONTS,
        scale=1.0,
    )

    trunk_action = JointPositionActionCfg(
        asset_name="robot",
        joint_names=[MINDBOT_TRUNK_JOINT],
        scale=1.0,
        use_default_offset=False,  # action = absolute target, not offset from default
    )


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

@configclass
class MindRobotDualArmObsCfg:
    """Observations for dual-arm teleoperation."""

    @configclass
    class PolicyCfg(ObsGroup):
        actions = ObsTerm(func=mdp.last_action)
        joint_pos = ObsTerm(func=mdp.joint_pos_rel)
        joint_vel = ObsTerm(func=mdp.joint_vel_rel)
        # Left arm EEF
        eef_pos_left = ObsTerm(func=mdp.ee_frame_pos)
        eef_quat_left = ObsTerm(func=mdp.ee_frame_quat)
        # Left gripper (2 joints)
        left_gripper_pos = ObsTerm(
            func=mdp.gripper_pos,
            params={"joint_names": MINDBOT_LEFT_GRIPPER_JOINTS},
        )
        # Right arm EEF
        eef_pos_right = ObsTerm(
            func=mdp.ee_frame_pos,
            params={"ee_frame_cfg": SceneEntityCfg("ee_frame_right")},
        )
        eef_quat_right = ObsTerm(
            func=mdp.ee_frame_quat,
            params={"ee_frame_cfg": SceneEntityCfg("ee_frame_right")},
        )
        # Right gripper (2 joints)
        right_gripper_pos = ObsTerm(
            func=mdp.gripper_pos,
            params={"joint_names": MINDBOT_RIGHT_GRIPPER_JOINTS},
        )

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

    policy: PolicyCfg = PolicyCfg()


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

@configclass
class MindRobotDualArmTerminationsCfg:
    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 at startup."""
    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)


def _disable_trunk_head_gravity(env, env_ids: torch.Tensor):
    """Disable gravity for trunk and head at startup."""
    import isaaclab.sim.schemas as schemas

    cfg = RigidBodyPropertiesCfg(disable_gravity=True)
    for env_id in range(env.num_envs):
        for prim_path in [
            f"/World/envs/env_{env_id}/Robot/robot_trunk",
            f"/World/envs/env_{env_id}/Robot/robot_head",
        ]:
            schemas.modify_rigid_body_properties(prim_path, cfg)


@configclass
class MindRobotDualArmEventsCfg:
    disable_arm_gravity = EventTerm(
        func=_disable_arm_gravity,
        mode="startup",
    )
    disable_trunk_head_gravity = EventTerm(
        func=_disable_trunk_head_gravity,
        mode="startup",
    )
    reset_scene = EventTerm(
        func=mdp.reset_scene_to_default,
        mode="reset",
    )


# =====================================================================
# Empty placeholder configs
# =====================================================================

@configclass
class EmptyCfg:
    pass


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

@configclass
class MindRobot2iDualArmIKAbsEnvCfg(ManagerBasedRLEnvCfg):
    """MindRobot 双臂 IK-Abs 遥操作环境配置（XR 控制器）。"""

    scene: MindRobotDualArmSceneCfg = MindRobotDualArmSceneCfg(
        num_envs=1,
        env_spacing=2.5,
    )

    observations: MindRobotDualArmObsCfg = MindRobotDualArmObsCfg()
    actions: MindRobotDualArmActionsCfg = MindRobotDualArmActionsCfg()
    terminations: MindRobotDualArmTerminationsCfg = MindRobotDualArmTerminationsCfg()
    events: MindRobotDualArmEventsCfg = MindRobotDualArmEventsCfg()

    commands: EmptyCfg = EmptyCfg()
    rewards: EmptyCfg = EmptyCfg()
    curriculum: EmptyCfg = EmptyCfg()

    xr: XrCfg = XrCfg(
        anchor_pos=(-0.1, -0.5, -1.05),
        anchor_rot=(0.866, 0, 0, -0.5),
    )

    def __post_init__(self):
        super().__post_init__()

        self.viewer = ViewerCfg(
            eye=(2.0, -1.5, 1.8),      # 相机位置
            lookat=(0.0, 0.0, 0.8),    # 瞄准点（机器人腰部高度附近）
            origin_type="world",
        )
        self.decimation = 2  # 25 Hz control (was 2 = 50 Hz, unnecessary for teleop)
        self.episode_length_s = 50.0
        self.sim.dt = 0.01  # 100 Hz
        self.sim.render_interval = 2

        # Mobile base — root link not fixed
        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

        # Left arm EEF frame transformer
        marker_cfg = FRAME_MARKER_CFG.copy()
        marker_cfg.markers["frame"].scale = (0.1, 0.1, 0.1)
        marker_cfg.prim_path = "/Visuals/FrameTransformerLeft"
        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]),
                ),
            ],
        )

        # Right arm EEF frame transformer
        marker_cfg_right = FRAME_MARKER_CFG.copy()
        marker_cfg_right.markers["frame"].scale = (0.1, 0.1, 0.1)
        marker_cfg_right.prim_path = "/Visuals/FrameTransformerRight"
        self.scene.ee_frame_right = FrameTransformerCfg(
            prim_path=f"{{ENV_REGEX_NS}}/Robot/{MINDBOT_RIGHT_ARM_PRIM_ROOT}/base_link",
            debug_vis=False,
            visualizer_cfg=marker_cfg_right,
            target_frames=[
                FrameTransformerCfg.FrameCfg(
                    prim_path=f"{{ENV_REGEX_NS}}/Robot/{MINDBOT_RIGHT_EEF_PRIM}",
                    name="end_effector",
                    offset=OffsetCfg(pos=[0.0, 0.0, 0.0]),
                ),
            ],
        )

        # 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 params (left gripper for single-arm compat; dual-arm uses obs directly)
        self.gripper_joint_names = MINDBOT_LEFT_GRIPPER_JOINTS
        self.gripper_open_val = 0.0
        self.gripper_threshold = 0.005
