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

"""评估协议预设。

协议就是一次评估的全部「口径」。默认预设 ``standard_v1`` 的数值与旧
``scripts/reinforcement_learning/rsl_rl/eval.py`` 的命令行默认值一致，
所以两套评估跑出来的数字可以直接对照。

要改口径，改预设或另加一个预设，不要在调用处临时传参——预设名会进
``protocol_fingerprint``，写进输出 JSON。
"""

from __future__ import annotations

from dataclasses import replace

from eval_kit.core.schema import EvalProtocol

#: 标准越障评估：10 s、前进 1 m/s、最多 2 m 算成功。与旧 eval.py 默认一致。
#:
#: **只建议用于与旧 ``eval.py`` 做对照**：它的地砖 2.0 m 会让越界终止项误判最外侧格子，
#: 起始高度 0.0 m 会让机器人在倒金字塔台阶地形里被生成到几何体内部
#: （``stairs_down`` 因此恒为 0）。要拿可用的数字请用 ``standard_v1_clean``。
STANDARD_V1 = EvalProtocol(
    steps=500,
    dt=0.02,
    command=(1.0, 0.0, 0.0),
    start_pose=(-0.9, 0.0, 0.0),
    success_distance_m=2.0,
    tilt_threshold_rad=0.70,
    acc_threshold_g=2.0 / 9.81,
    energy_convention="absolute",
    stability_margin_method="foot_bbox",
    support_force_threshold_n=15.0,
    enable_interval_pushes=False,
    terrain_spec="standard_v1",
    terrain_size_m=2.0,
    num_levels=8,
    envs_per_cell=1,
    seeds=(0,),
)

#: 抗扰评估：每 2.01 s 施加一次水平速度扰动，其余同 standard_v1。
ROBUSTNESS_V1 = replace(
    STANDARD_V1,
    enable_interval_pushes=True,
    push_interval_s=(2.0, 2.01),
    push_velocity_range=(-0.5, 0.5),
)

#: 冒烟测试：步数与网格都缩小，用于确认流水线能跑通，不要用它的数值下结论。
#: 地形参数沿用 ``standard_v1_clean`` 的修正值，否则它会带上越界误判的假警告。
SMOKE = replace(
    STANDARD_V1,
    steps=100,
    num_levels=2,
    envs_per_cell=1,
    terrain_size_m=4.0,
    start_z_offset_m=0.5,
)

#: **推荐用于机器人对比**：修掉两处会让结果失真的问题，其余口径与 ``standard_v1`` 相同，
#: 因此数值仍可与后者对照。
#:
#: 1. 地砖从 2.0 m 放大到 4.0 m，消除 IsaacLab ``terrain_out_of_bounds`` 的误判
#:    （阈值 3.8 m 由 ``EvalProtocol.minimum_terrain_size_for_bounds()`` 给出）。
#: 2. 起始高度抬高 0.5 m，让机器人落到地形表面而不是被生成在地形内部。
#:    实测证据：``--terrain-type stairs_down`` 单独跑，起始高度 0.0 m 时
#:    score=0.00 / 成功率 0.05，抬高后 score 显著恢复。取 0.5 是因为倒金字塔台阶的
#:    嵌入深度随难度递增（L1 约 0.11 m、L8 约 0.50 m），0.5 m 刚好覆盖最深的一级。
#:    旧 ``eval.py`` 用 0.0，所以它的倒金字塔台阶那一列恒为 0，是无效的。
STANDARD_V1_CLEAN = replace(STANDARD_V1, terrain_size_m=4.0, start_z_offset_m=0.5)

#: 另一种处置方式：保留 2.0 m 地砖，直接去掉越界终止项。
STANDARD_V1_NO_BOUNDS = replace(STANDARD_V1, enforce_terrain_bounds=False)

#: 多随机种子的标准评估 —— **对外给结论时用这个**。
#:
#: 单次 rollout 的噪声足以淹没机器人之间的差异，所以对比必须多种子；同时它基于
#: ``STANDARD_V1_CLEAN``（地砖 4.0 m、起始高度 +0.5 m），不带已知的 harness 偏差。
#:
#: 默认 ``envs_per_cell = 1``（每格 1 个机器人），因为起始位姿钉死后同格多个机器人
#: 起始坐标完全相同、不是独立样本。样本量靠种子数而非每格机器人数量来提升 ——
#: 每个种子会重新生成地形实例，跨种子才是独立的。8 个种子给出每格 1/8 ≈ 12.5% 的
#: 分辨率；想要更细就显式加 ``--seeds``。
STANDARD_V1_MULTISEED = replace(STANDARD_V1_CLEAN, seeds=(0, 100, 200, 300, 400, 500, 600, 700))

PROTOCOLS: dict[str, EvalProtocol] = {
    "standard_v1": STANDARD_V1,
    "standard_v1_clean": STANDARD_V1_CLEAN,
    "standard_v1_no_bounds": STANDARD_V1_NO_BOUNDS,
    "robustness_v1": ROBUSTNESS_V1,
    "smoke": SMOKE,
    "standard_v1_multiseed": STANDARD_V1_MULTISEED,
}


def get_protocol(name: str) -> EvalProtocol:
    if name not in PROTOCOLS:
        raise KeyError(f"未知评估协议 {name!r}；可选：{sorted(PROTOCOLS)}")
    return PROTOCOLS[name]