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

"""指标计算：纯函数，只依赖 numpy。

输入是 :class:`~eval_kit.core.schema.Trace`、:class:`~eval_kit.core.schema.RobotProfile`
和 :class:`~eval_kit.core.schema.EvalProtocol`，输出是一个扁平的指标字典。

指标口径与 ``scripts/reinforcement_learning/rsl_rl/eval.py`` 保持一致（同样的取点、
同样的公式），因此默认协议下两套评估的数值可以直接对照；差别只在三处，且都是
有意为之：

1. **质量**不再写死在命令行默认值里，而是从 ``Trace.mass_kg`` 取（适配层从资产读），
   传入的 ``profile.mass_kg`` 只作兜底。
2. **加速度阈值**以重力倍数表达，换算成 m/s² 后参与计算，使阈值可跨量级机器人复用。
3. **缺失的形态**返回 ``nan`` 而不是 0。旧实现里没有足端的机器人会得到
   「支撑裕度 = 0.00」，看起来像测出来的值，实际是没测。
"""

from __future__ import annotations

from typing import Any

import numpy as np

from eval_kit.core.schema import EvalProtocol, RobotProfile, Trace

#: 全部帧，用于成功判定与最大进度（含起始帧，与旧实现取点一致）。
_ALL_FRAMES = slice(None)
#: 动作执行之后的帧，用于逐帧累积类指标。
_STEP_FRAMES = slice(1, None)

_GRAVITY = 9.81


def tilt_from_projected_gravity(projected_gravity_b: np.ndarray) -> np.ndarray:
    """由机体系投影重力算倾角 [rad]。

    ``projected_gravity_b`` 为 ``(T, N, 3)``，返回 ``(T, N)``。
    """
    z_axis = np.clip(-projected_gravity_b[..., 2], -1.0, 1.0)
    return np.arccos(z_axis)


def support_margin_xy(root_xy: np.ndarray, foot_xy: np.ndarray) -> np.ndarray:
    """机身投影点到足端轴对齐包围盒的距离 [m]，在盒内为正、盒外为负。

    ``root_xy`` 为 ``(N, 2)``，``foot_xy`` 为 ``(N, F, 2)``，返回 ``(N,)``。

    这是当前实现采用的口径（与旧 ``eval.py`` 相同），记录在
    ``EvalProtocol.stability_margin_method`` 里。它比凸包/ZMP 版本粗糙，
    但不会在退化多边形上失败，也不需要 scipy。
    """
    min_xy = np.min(foot_xy, axis=1)
    max_xy = np.max(foot_xy, axis=1)

    inside_x = (root_xy[:, 0] >= min_xy[:, 0]) & (root_xy[:, 0] <= max_xy[:, 0])
    inside_y = (root_xy[:, 1] >= min_xy[:, 1]) & (root_xy[:, 1] <= max_xy[:, 1])
    inside = inside_x & inside_y

    inside_dist = np.minimum(
        np.minimum(root_xy[:, 0] - min_xy[:, 0], max_xy[:, 0] - root_xy[:, 0]),
        np.minimum(root_xy[:, 1] - min_xy[:, 1], max_xy[:, 1] - root_xy[:, 1]),
    )
    outside_dx = np.clip(np.maximum(min_xy[:, 0] - root_xy[:, 0], root_xy[:, 0] - max_xy[:, 0]), 0.0, None)
    outside_dy = np.clip(np.maximum(min_xy[:, 1] - root_xy[:, 1], root_xy[:, 1] - max_xy[:, 1]), 0.0, None)
    outside_dist = np.sqrt(outside_dx**2 + outside_dy**2)

    return np.where(inside, inside_dist, -outside_dist)


def _nanmean_or_nan(values: np.ndarray) -> float:
    """均值；空数组返回 nan（而不是 0，也不是 numpy 的警告）。"""
    if values.size == 0:
        return float("nan")
    return float(np.mean(values))


def compute_metrics(
    trace: Trace,
    profile: RobotProfile,
    protocol: EvalProtocol,
    *,
    include_per_env: bool = False,
) -> dict[str, Any]:
    """由一次 rollout 的轨迹算出全部指标。

    指标名的命名与旧 ``eval.py`` 完全一致，便于新旧结果直接对照。
    """
    num_steps = trace.num_steps
    num_envs = trace.num_envs
    dt = float(trace.dt)

    if abs(protocol.dt - dt) > 0.01 * dt:
        raise ValueError(
            f"协议 dt={protocol.dt} 与轨迹 dt={dt} 不一致。"
            "通常意味着协议里的步长与任务的 decimation 不匹配。"
        )

    command = np.asarray(trace.command, dtype=np.float64)
    command_xy = command[:2]
    command_yaw = float(command[2])

    command_dir = command_xy.copy()
    norm = float(np.linalg.norm(command_dir))
    command_dir = command_dir / norm if norm > 1e-6 else np.array([1.0, 0.0])

    # ------------------------------------------------------------------ 成功判定
    base_xy = np.asarray(trace.base_pos_w, dtype=np.float64)[..., :2]
    delta_xy = base_xy - base_xy[0][None, :, :]
    progress = np.einsum("tnc,c->tn", delta_xy[_ALL_FRAMES], command_dir)
    max_progress = progress.max(axis=0)
    reached = max_progress >= protocol.success_distance_m

    tilt = tilt_from_projected_gravity(np.asarray(trace.projected_gravity_b, dtype=np.float64))
    tilt_steps = tilt[_STEP_FRAMES]
    max_tilt = tilt_steps.max(axis=0)
    ever_fallen = max_tilt > protocol.tilt_threshold_rad

    success = reached & ~ever_fallen
    success_count = int(np.count_nonzero(success))

    # 首次达成成功的时刻：第 i 帧（i>=1）达成则记 i*dt。
    crossed = progress[_STEP_FRAMES] >= protocol.success_distance_m
    complete_time = np.full((num_envs,), np.nan, dtype=np.float64)
    for env_index in range(num_envs):
        hit = np.flatnonzero(crossed[:, env_index])
        if hit.size:
            complete_time[env_index] = float(hit[0] + 1) * dt

    successful_time = complete_time[success & np.isfinite(complete_time)]
    avg_time = _nanmean_or_nan(successful_time)

    # ------------------------------------------------------------- 速度跟踪误差
    lin_vel_b = np.asarray(trace.base_lin_vel_b, dtype=np.float64)[_STEP_FRAMES]
    ang_vel_b = np.asarray(trace.base_ang_vel_b, dtype=np.float64)[_STEP_FRAMES]
    lin_error = np.linalg.norm(lin_vel_b[..., :2] - command_xy[None, None, :], axis=-1)
    yaw_error = np.abs(ang_vel_b[..., 2] - command_yaw)

    # -------------------------------------------------------------- 能耗与路程
    joint_torque = np.asarray(trace.joint_torque, dtype=np.float64)[_STEP_FRAMES]
    joint_vel = np.asarray(trace.joint_vel, dtype=np.float64)[_STEP_FRAMES]
    mechanical_power = joint_torque * joint_vel
    if protocol.energy_convention == "absolute":
        integrand = np.abs(mechanical_power)
    elif protocol.energy_convention == "signed":
        integrand = mechanical_power
    else:
        raise ValueError(
            f"未知能耗约定 {protocol.energy_convention!r}；可选 'absolute' 或 'signed'。"
        )
    energy = integrand.sum(axis=(0, 2)) * dt

    lin_vel_w = np.asarray(trace.base_lin_vel_w, dtype=np.float64)
    route = np.linalg.norm(lin_vel_w[_STEP_FRAMES][..., :2], axis=-1).sum(axis=0) * dt

    # ------------------------------------------------------------ 运动不稳定性
    accel = np.linalg.norm(np.diff(lin_vel_w, axis=0), axis=-1) / dt
    instability_count = (accel > protocol.acc_threshold_mps2).sum(axis=0)

    stable_time = (tilt_steps < protocol.tilt_threshold_rad).sum(axis=0) * dt

    # ------------------------------------------------------------- 支撑多边形
    feet_pos_w = np.asarray(trace.feet_pos_w, dtype=np.float64)[_STEP_FRAMES]
    num_feet = int(feet_pos_w.shape[2])
    if num_feet > 0:
        margins = support_margin_xy(base_xy[_STEP_FRAMES][:, :, :].reshape(-1, 2), feet_pos_w.reshape(-1, num_feet, 2))
        margins = margins.reshape(feet_pos_w.shape[0], num_envs)
        avg_margin = _nanmean_or_nan(margins)
        margin_method = protocol.stability_margin_method
    else:
        avg_margin = float("nan")
        margin_method = "unavailable"

    # --------------------------------------------------------------- 支撑接触
    support_force = np.asarray(trace.support_force_n, dtype=np.float64)
    support_contact_time = np.asarray(trace.support_contact_time_s, dtype=np.float64)
    num_support_bodies = int(support_force.shape[2])
    support_threshold = (
        profile.support_force_threshold_n
        if protocol.support_force_threshold_n is None
        else protocol.support_force_threshold_n
    )
    if num_support_bodies > 0:
        active = np.any(support_force[_STEP_FRAMES] > support_threshold, axis=2)  # (T-1, N)
        support_steps = active.sum(axis=0)
        # 起始帧之前视为「未接触」，因此首帧接触也算一次事件（与旧实现一致）。
        active_padded = np.concatenate([np.zeros((1, num_envs), dtype=bool), active], axis=0)
        support_events = (active_padded[1:] & ~active_padded[:-1]).sum(axis=0)
        max_support_time = support_contact_time[_STEP_FRAMES].max(axis=(0, 2))
        support_contact_ratio = _nanmean_or_nan(support_steps / max(num_steps, 1))
        avg_support_events = float(np.mean(support_events))
        avg_max_support_time = float(np.mean(max_support_time))
        if success_count > 0:
            success_support_ratio = float(np.mean((support_events[success] > 0).astype(np.float64)))
        else:
            success_support_ratio = float("nan")
    else:
        support_contact_ratio = float("nan")
        avg_support_events = float("nan")
        avg_max_support_time = float("nan")
        success_support_ratio = float("nan")

    # --------------------------------------------------------- 质量与 COT
    mass = trace.mass_kg if trace.mass_kg is not None else profile.mass_kg
    if mass is None or mass <= 0.0:
        cost_of_transport = np.full((num_envs,), np.nan)
        mass_used = float("nan")
    else:
        cost_of_transport = energy / ((route + 1e-6) * mass * _GRAVITY)
        mass_used = float(mass)

    metrics: dict[str, Any] = {
        "profile": profile.name,
        "protocol_fingerprint": protocol.fingerprint(),
        "num_envs": float(num_envs),
        "num_steps": float(num_steps),
        "episode_length_s": float(num_steps * dt),
        "terrain_size_m": float(protocol.terrain_size_m),
        "success_distance_m": float(protocol.success_distance_m),
        "robot_mass_kg": mass_used,
        "energy_convention": protocol.energy_convention,
        "stability_margin_method": margin_method,
        "tilt_threshold_rad": float(protocol.tilt_threshold_rad),
        "acc_threshold_mps2": float(protocol.acc_threshold_mps2),
        "support_force_threshold_n": float(support_threshold) if num_support_bodies > 0 else float("nan"),
        # --- 成功率与进度
        "success_count": float(success_count),
        "success_rate": float(success_count / max(num_envs, 1)),
        "avg_final_progress_m": float(np.mean(max_progress)),
        # --- 时间
        "avg_time_consumption_s": avg_time,
        "time_efficiency": float(
            avg_time / (protocol.success_distance_m / max(abs(float(command[0])), 1e-6))
        )
        if np.isfinite(avg_time)
        else float("nan"),
        # --- 跟踪误差
        "mean_lin_vel_error": float(np.mean(lin_error)),
        "mean_yaw_vel_error": float(np.mean(yaw_error)),
        # --- 能耗
        "avg_energy_j": float(np.mean(energy)),
        "avg_cost_of_transport": _nanmean_or_nan(cost_of_transport),
        # --- 稳定性
        "unstability_rate": float(np.mean(instability_count / max(num_steps, 1))),
        "avg_stability_margin_m": avg_margin,
        "avg_max_tilt_rad": float(np.mean(max_tilt)),
        "avg_stable_time_s": float(np.mean(stable_time)),
        # --- 支撑接触
        "avg_support_contact_ratio": support_contact_ratio,
        "avg_support_events": avg_support_events,
        "avg_max_support_contact_time_s": avg_max_support_time,
        "successful_runs_with_support_ratio": success_support_ratio,
    }

    if include_per_env:
        metrics["_per_env"] = {
            "success": success.astype(float).tolist(),
            "max_progress": max_progress.tolist(),
            "max_tilt": max_tilt.tolist(),
        }
    return metrics