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

"""评价套件的数据契约。

三个对象构成整条流水线：

``RobotProfile``   机器人形态（谁是足端、谁是支撑体、质量多少）
``EvalProtocol``   评估协议（指令、时长、各种阈值、能耗约定）
``Trace``          一次 rollout 的标准化逐帧记录

上游（任何框架的适配器）负责产出 ``Trace``，下游（:mod:`eval_kit.core.metrics`）
只从 ``Trace`` 计算指标。指标代码因此不接触任何仿真对象。

这个模块只依赖 numpy 与标准库。
"""

from __future__ import annotations

import hashlib
import json
from dataclasses import asdict, dataclass
from typing import Any

import numpy as np

#: 帧索引约定：``Trace`` 的第 0 帧是执行动作之前的起始状态，
#: 第 ``i`` 帧（``i >= 1``）是第 ``i`` 步动作之后的状态。
#: 因此一次 ``steps`` 步的 rollout 会产生 ``steps + 1`` 帧。
START_FRAME = 0

#: 成功判定与最大进度使用全部帧（含起始帧），与 legged_gym 侧
#: ``criterion.py`` 及现有 ``eval.py`` 的取点一致。
PROGRESS_FRAMES = slice(None)

#: 逐帧累积类指标只使用动作之后的帧。
STEP_FRAMES = slice(1, None)


@dataclass(frozen=True)
class RobotProfile:
    """机器人形态描述。

    评价代码里不应出现具体的 link/joint 名字或质量常数，全部集中在这里。
    新增一个机器人 = 新增一条 ``RobotProfile`` 记录，不改评价代码。
    """

    name: str
    """机器人标识，同时也是输出文件名的一部分。"""

    mass_kg: float | None
    """整机质量 [kg]；``None`` 表示运行时从仿真资产读取（推荐）。"""

    feet_body_patterns: tuple[str, ...] = ()
    """支撑多边形的接触点，例如火蚁的四个轮。空元组表示该机器人不计算支撑裕度。"""

    support_body_patterns: tuple[str, ...] = ()
    """允许作为「短暂支撑」的机体，例如火蚁的机身与膝部。空元组表示不做支撑统计。"""

    actuated_joint_patterns: tuple[str, ...] = ()
    """参与能耗统计的关节。"""

    support_force_threshold_n: float = 15.0
    """判定「支撑接触」的法向力阈值 [N]。"""

    body_length_m: float = 0.0
    """机身长度 [m]，仅用于校验地形瓦片尺寸是否够大。"""

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass(frozen=True)
class EvalProtocol:
    """评估协议：一次评估的全部「口径」。

    这些字段决定数值大小，因此都会被写进输出 JSON 的 ``fingerprint``，
    避免不同设置下的结果被悄悄拿来比较。

    默认值与现有 ``scripts/reinforcement_learning/rsl_rl/eval.py`` 的默认
    命令行参数一致（能耗用绝对值口径、倾角阈值 0.70 rad、加速度阈值 2.0 m/s²
    折算成重力倍数），因此默认协议下的新结果可以直接和旧评估结果对照。
    """

    steps: int = 500
    """每个 episode 的控制步数。dt=0.02 时 500 步 = 10 s。"""

    dt: float = 0.02
    """控制步长 [s]，即 ``decimation * sim.dt``。"""

    command: tuple[float, float, float] = (1.0, 0.0, 0.0)
    """固定速度指令 ``(lin_vel_x, lin_vel_y, ang_vel_z)``。"""

    start_pose: tuple[float, float, float] = (-0.9, 0.0, 0.0)
    """起始位姿 ``(x, y, yaw)``，相对地形瓦片中心。"""

    start_z_offset_m: float = 0.0
    """起始高度相对地形瓦片中心高度的额外抬升 [m]。

    默认 0.0 表示把机器人直接放在 ``env_origin.z`` 上，与旧 ``eval.py`` 的写法一致。
    但那样只对「瓦片中心是最高点」的地形成立：像倒金字塔台阶（``pyramid_stairs_inv``，
    中心是最低点）这类地形，机器人在 ``|x| > platform`` 处会被直接生成到地形几何体
    内部，卡住无法移动，导致该地形在所有难度上都得到 0 成功率。

    给一个正值（例如 0.3 m）让机器人从地形表面上方落下，可以适配任意地形形状。
    改这个值会改变 ``protocol_fingerprint``。
    """

    success_distance_m: float = 2.0
    """沿指令方向前进多少米算成功。"""

    tilt_threshold_rad: float = 0.70
    """倾角超过该值 [rad] 视为摔倒。0.70 rad ≈ 40°。"""

    acc_threshold_g: float = 2.0 / 9.81
    """机身加速度超过该值（单位 g）计为一次不稳定。用重力倍数而非绝对 m/s²，
    使该阈值对不同量级的机器人保持可比。默认值 2.0 m/s² 折算而来。"""

    energy_convention: str = "absolute"
    """能耗积分约定，``"absolute"`` 表示 ``Σ|τ·q̇|·dt``，
    ``"signed"`` 表示 ``Σ(τ·q̇)·dt``。两者数值差异可达一倍以上，
    全局必须一致；该值会写进输出 JSON。"""

    stability_margin_method: str = "foot_bbox"
    """支撑裕度算法。``"foot_bbox"`` = 机体到足端轴对齐包围盒的距离（当前实现）。"""

    support_force_threshold_n: float | None = 15.0
    """判定支撑接触的法向力阈值 [N]；``None`` 时改用机器人描述里的值。"""

    enable_interval_pushes: bool = False
    """是否启用周期性外力扰动（``robustness`` 评估用）。"""

    enforce_terrain_bounds: bool = True
    """是否保留 IsaacLab 的 ``terrain_out_of_bounds`` 终止项。

    **默认 True 是为了与旧 ``eval.py`` 的结果可比，但已知它有系统性偏差**：
    该终止项比较的是机器人的世界坐标与整张地图的范围，而不是它自己那块地砖。
    在地砖 2 m、网格 5x5 或 8xN 的评估配置下，最外侧的格子会被误判为越界，
    机器人从第 0 步起每步都被重置，永远没机会尝试地形；而 ``|y|`` 恰好等于阈值
    的那两列则处于浮点边界上。实测：8x2 网格下首步就有 6/32 个环境被误判。

    三种处置方式（建议前两种，按顺序）：

    1. 把 :attr:`terrain_size_m` 提到 **≥ 3.8**，即可让所有起始位置落进阈值内，
       同时保留这个终止项；
    2. 置为 ``False``，评估时去掉这个终止项（成功判定本来就是按前进距离算的，
       走出地砖不构成失败）；
    3. 保持 ``True`` 且用小地砖 —— 数值会偏低，且被影响的难度等级因行序打乱而
       每次不同，不建议用来做对比。
    """

    push_interval_s: tuple[float, float] = (2.0, 2.01)
    """扰动间隔 [s]，取一个极窄区间让扰动近似等周期但可比较。"""

    push_velocity_range: tuple[float, float] = (-0.5, 0.5)
    """每个水平方向上的扰动速度范围 [m/s]。"""

    terrain_spec: str = "standard_v1"
    """地形等级规格名，见 :mod:`eval_kit.core.terrain_spec`。"""

    terrain_size_m: float = 2.0
    """单块地形瓦片边长 [m]。"""

    num_levels: int = 8
    """上限评估的难度等级数（网格行数）。"""

    envs_per_cell: int = 1
    """上限评估每个「地形类型 × 难度」格子里放几个机器人。

    默认 1。**注意：由于起始位姿被钉死，同一格的多个机器人起始坐标完全相同**，
    它们不是独立样本，差异只来自物理求解器的非确定性。所以靠「一格放多个」来提高
    样本量是假精度；正确做法是 ``envs_per_cell=1`` + 多个 :attr:`seeds`
    （不同种子会生成不同的地形实例，因此跨种子才是独立样本）。

    代价：``envs_per_cell=1`` 时单次评估里每格成功率只能是 0 或 1，
    必须配合足够多的种子才能得到有意义的分辨率。经验值：种子数 ≥ 8 时
    每格分辨率达到 1/8 ≈ 12.5%。
    """

    seeds: tuple[int, ...] = (0,)
    """重复试验的随机种子。单次 rollout 噪声很大，比较机器人时应跑多个种子。

    每个种子会重新生成地形实例（``terrain_generator.seed = seed``），所以跨种子
    的样本是独立的；这是提高统计可靠性的正确途径，而不是加大 ``envs_per_cell``。
    """

    @property
    def acc_threshold_mps2(self) -> float:
        """加速度阈值的 m/s² 形式，供指标计算使用。"""
        return self.acc_threshold_g * 9.81

    def minimum_terrain_size_for_bounds(self, *, border_width: float = 2.0, distance_buffer: float = 3.0) -> float:
        """保留 ``terrain_out_of_bounds`` 终止项且不误判所需的最小地砖边长 [m]。

        这个终止项判定 ``|origin + start| <= 0.5 * 地图边长 - distance_buffer``，
        而最外侧格子中心在 ``(rows-1)/2 * tile``，地图边长 ``= rows * tile + 2 * border``。
        两边的 ``rows`` 项相消，于是结论与网格行数无关：

            ``tile >= 2 * (|start_x| + distance_buffer - border_width)``

        默认的 ``border_width=2.0`` / ``distance_buffer=3.0`` 与
        :func:`eval_kit.isaaclab.env_setup.apply_eval_terrain_baseline` 及 IsaacLab 的
        默认值一致；取 ``start_x = -0.9`` 得 ``tile >= 3.8 m``。

        放在内核里（而不是适配层）是为了让编排进程也能用——它不启动仿真，
        无法 import 任何依赖 isaaclab 的模块。
        """
        return max(0.0, 2.0 * (abs(self.start_pose[0]) + distance_buffer - border_width))

    def fingerprint(self) -> str:
        """协议指纹：决定数值的字段的稳定哈希，写进输出 JSON。"""
        payload = {
            "steps": self.steps,
            "dt": self.dt,
            "command": list(self.command),
            "start_pose": list(self.start_pose),
            "start_z_offset_m": self.start_z_offset_m,
            "success_distance_m": self.success_distance_m,
            "tilt_threshold_rad": self.tilt_threshold_rad,
            "acc_threshold_g": self.acc_threshold_g,
            "energy_convention": self.energy_convention,
            "stability_margin_method": self.stability_margin_method,
            "support_force_threshold_n": self.support_force_threshold_n,
            "enable_interval_pushes": self.enable_interval_pushes,
            "enforce_terrain_bounds": self.enforce_terrain_bounds,
            "push_interval_s": list(self.push_interval_s),
            "push_velocity_range": list(self.push_velocity_range),
            "terrain_spec": self.terrain_spec,
            "terrain_size_m": self.terrain_size_m,
            "num_levels": self.num_levels,
            "envs_per_cell": self.envs_per_cell,
        }
        blob = json.dumps(payload, sort_keys=True).encode("utf-8")
        return hashlib.sha256(blob).hexdigest()[:12]

    def to_dict(self) -> dict[str, Any]:
        data = asdict(self)
        data["acc_threshold_mps2"] = self.acc_threshold_mps2
        data["fingerprint"] = self.fingerprint()
        return data


@dataclass
class Trace:
    """一次 rollout 的标准化逐帧记录。

    所有数组的第一维都是帧、第二维都是环境。第 0 帧是动作执行之前的起始状态。

    ``feet_pos_w`` / ``support_force_n`` / ``support_contact_time_s`` 的第二维
    可能是 0（该机器人没有足端或没有支撑体语义），指标层会跳过对应指标。
    """

    dt: float
    command: np.ndarray
    """``(3,)`` 本次 rollout 使用的固定指令。"""

    base_pos_w: np.ndarray
    """``(T, N, 3)`` 机身世界位置。"""

    base_lin_vel_w: np.ndarray
    """``(T, N, 3)`` 机身世界系线速度，用于路程与加速度（不稳定性）。"""

    base_lin_vel_b: np.ndarray
    """``(T, N, 3)`` 机体系线速度，用于速度跟踪误差。"""

    base_ang_vel_b: np.ndarray
    """``(T, N, 3)`` 机体系角速度，用第 2 个分量做 yaw 跟踪误差。"""

    projected_gravity_b: np.ndarray
    """``(T, N, 3)`` 机体系投影重力，用于倾角。"""

    joint_torque: np.ndarray
    """``(T, N, J)`` 实际施加的关节力矩。"""

    joint_vel: np.ndarray
    """``(T, N, J)`` 关节速度。"""

    feet_pos_w: np.ndarray
    """``(T, N, F, 2)`` 足端/轮心的世界 xy，用于支撑多边形。"""

    support_force_n: np.ndarray
    """``(T, N, S)`` 支撑体的法向合力。"""

    support_contact_time_s: np.ndarray
    """``(T, N, S)`` 支撑体的连续接触时长。"""

    done: np.ndarray
    """``(T, N)`` 每步之后的终止标志（1.0 表示该步结束后被重置）。"""

    mass_kg: float | None = None
    """实际使用的整机质量 [kg]。``None`` 表示上游未能从资产读到。"""

    def __post_init__(self) -> None:
        self.validate()

    @property
    def num_frames(self) -> int:
        return int(self.base_pos_w.shape[0])

    @property
    def num_steps(self) -> int:
        """实际执行的控制步数（帧数减一）。"""
        return self.num_frames - 1

    @property
    def num_envs(self) -> int:
        return int(self.base_pos_w.shape[1])

    @property
    def num_joints(self) -> int:
        return int(self.joint_torque.shape[2])

    def validate(self) -> None:
        """形状一致性检查。早失败好过算出一堆形状广播出来的错数字。"""
        frames, envs = self.num_frames, self.num_envs
        if frames < 2:
            raise ValueError(f"Trace 至少需要 2 帧（起始帧 + 至少一步），收到 {frames} 帧。")

        expected: dict[str, tuple[int, ...]] = {
            "base_pos_w": (frames, envs, 3),
            "base_lin_vel_w": (frames, envs, 3),
            "base_lin_vel_b": (frames, envs, 3),
            "base_ang_vel_b": (frames, envs, 3),
            "projected_gravity_b": (frames, envs, 3),
            "done": (frames, envs),
        }
        for name, shape in expected.items():
            array = np.asarray(getattr(self, name))
            if array.shape != shape:
                raise ValueError(f"{name} 形状应为 {shape}，实际 {array.shape}。")

        joints = self.num_joints
        for name in ("joint_torque", "joint_vel"):
            if np.asarray(getattr(self, name)).shape != (frames, envs, joints):
                raise ValueError(
                    f"{name} 形状应为 {(frames, envs, joints)}，实际 {np.asarray(getattr(self, name)).shape}。"
                )

        feet = np.asarray(self.feet_pos_w)
        if feet.shape[:2] != (frames, envs) or feet.ndim != 4 or feet.shape[3] != 2:
            raise ValueError(f"feet_pos_w 形状应为 {(frames, envs, feet.shape[2] if feet.ndim > 2 else 0, 2)}，"
                             f"实际 {feet.shape}。")

        for name in ("support_force_n", "support_contact_time_s"):
            array = np.asarray(getattr(self, name))
            if array.ndim != 3 or array.shape[:2] != (frames, envs):
                raise ValueError(f"{name} 形状应为 {(frames, envs, array.shape[2] if array.ndim > 2 else 0)}，"
                                 f"实际 {array.shape}。")

        if np.asarray(self.command).shape != (3,):
            raise ValueError(f"command 应为 (3,)，实际 {np.asarray(self.command).shape}。")

    def save(self, path: str) -> None:
        """把轨迹存成 ``.npz``，便于事后用新指标重新分析，无需重跑仿真。"""
        np.savez_compressed(
            path,
            dt=np.asarray(self.dt),
            command=self.command,
            base_pos_w=self.base_pos_w,
            base_lin_vel_w=self.base_lin_vel_w,
            base_lin_vel_b=self.base_lin_vel_b,
            base_ang_vel_b=self.base_ang_vel_b,
            projected_gravity_b=self.projected_gravity_b,
            joint_torque=self.joint_torque,
            joint_vel=self.joint_vel,
            feet_pos_w=self.feet_pos_w,
            support_force_n=self.support_force_n,
            support_contact_time_s=self.support_contact_time_s,
            done=self.done,
            mass_kg=np.asarray(-1.0 if self.mass_kg is None else self.mass_kg),
        )

    @classmethod
    def load(cls, path: str) -> Trace:
        with np.load(path) as data:
            mass = float(data["mass_kg"])
            return cls(
                dt=float(data["dt"]),
                command=data["command"],
                base_pos_w=data["base_pos_w"],
                base_lin_vel_w=data["base_lin_vel_w"],
                base_lin_vel_b=data["base_lin_vel_b"],
                base_ang_vel_b=data["base_ang_vel_b"],
                projected_gravity_b=data["projected_gravity_b"],
                joint_torque=data["joint_torque"],
                joint_vel=data["joint_vel"],
                feet_pos_w=data["feet_pos_w"],
                support_force_n=data["support_force_n"],
                support_contact_time_s=data["support_contact_time_s"],
                done=data["done"],
                mass_kg=None if mass < 0.0 else mass,
            )