# Copyright (c) 2025 Deep Robotics
# SPDX-License-Identifier: BSD 3-Clause

"""机器人形态注册表。

**加一个机器人 = 在这里加一条记录**，评价代码不动。

数值来源：火蚁的 link/joint 取自
``source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/rough_env_cfg.py``；
机身尺寸取自 ``docs/solenopsis_urdf_comparison.md`` 里记录的碰撞体尺寸。
"""

from __future__ import annotations

from eval_kit.core.schema import RobotProfile

SOLENOPSIS = RobotProfile(
    name="solenopsis",
    # None = 运行时从仿真资产读真实质量（Solenopsis.urdf 当前为 26.581 kg）。
    # 旧 eval.py 的命令行默认值是 25.482，那是更早的 M20 简化模型名义质量，
    # 会让 COT 偏低约 4%。
    mass_kg=None,
    # 四个轮就是支撑多边形的接触点。
    feet_body_patterns=(".*_wheel",),
    # 机身与膝部允许作为越障时的短暂支撑。
    support_body_patterns=("base_link", ".*_knee"),
    # 火蚁 8 个关节全部主动驱动；分两组写便于将来出现被动关节时收窄范围。
    actuated_joint_patterns=(".*knee_joint", ".*wheel_joint"),
    support_force_threshold_n=15.0,
    # 机身碰撞体 0.706 x 0.310 x 0.1714 m，取长边。
    body_length_m=0.706,
)

#: 四足对照机，仅用于验证「换机器人只改数据」这一点，尚未在仿真里跑过。
LITE3 = RobotProfile(
    name="lite3",
    mass_kg=None,
    feet_body_patterns=(".*_FOOT",),
    # 四足没有「机身支撑」这个语义，留空 → 支撑类指标返回 nan 而不是 0。
    support_body_patterns=(),
    actuated_joint_patterns=(".*_joint",),
    support_force_threshold_n=15.0,
    # 未知，仅用于地形瓦片尺寸校验。
    body_length_m=0.0,
)

PROFILES: dict[str, RobotProfile] = {
    SOLENOPSIS.name: SOLENOPSIS,
    LITE3.name: LITE3,
}


def get_profile(name: str) -> RobotProfile:
    if name not in PROFILES:
        raise KeyError(f"未知机器人描述 {name!r}；可选：{sorted(PROFILES)}")
    return PROFILES[name]