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

"""上限评估的网格汇总：把逐环境结果折叠成「难度等级 × 地形类型」的成功率矩阵。

纯数据变换，只依赖 numpy，与仿真框架无关。

``score`` 的公式与 ``criterion.py`` / 旧 ``eval.py`` 完全一致（三者代数等价）：
对每个地形类型，先算逐级累积成功率 ``acc[i] = Π_{j<=i} rate[j]``，
再算 ``score = Σ_i acc[i] * (i+1) * (1 - rate[i+1])``，最后一级的 ``(1 - rate)`` 取 1。
含义是「在爬到第 i+1 级之前失败」的累积期望，值越大说明越能爬到高难度。
"""

from __future__ import annotations

from typing import Any, Sequence

import numpy as np

from eval_kit.core.terrain_spec import TerrainSpec


def accumulated_success(rates: Sequence[float]) -> list[float]:
    """逐级累积成功率。"""
    accumulated: list[float] = []
    running = 1.0
    for rate in rates:
        running *= float(rate)
        accumulated.append(running)
    return accumulated


def upper_limit_score(rates: Sequence[float]) -> float:
    """由某一地形类型各难度等级的成功率算 score。"""
    rates = [float(rate) for rate in rates]
    accumulated = accumulated_success(rates)
    score = 0.0
    for index, acc_rate in enumerate(accumulated):
        next_failure = 1.0 - rates[index + 1] if index + 1 < len(rates) else 1.0
        score += acc_rate * (index + 1) * next_failure
    return score


def aggregate_runs(runs: Sequence[tuple[int, dict[str, Any]]]) -> dict[str, Any]:
    """把多次（不同随机种子）的指标聚合成 mean/std。

    单次 rollout 的噪声足以淹没两个策略之间的真实差异，所以做对比时应该跑多个
    种子并看分布，而不是看单个数。

    Args:
        runs: ``(seed, metrics)`` 列表，``metrics`` 来自
            :func:`eval_kit.core.metrics.compute_metrics`。
    """
    if not runs:
        raise ValueError("没有任何一次 run 可以聚合。")

    # 逐环境的私有字段不参与聚合
    metric_keys: list[str] = []
    for _, metrics in runs:
        for key, value in metrics.items():
            if key.startswith("_"):
                continue
            if isinstance(value, (int, float)) and key not in metric_keys:
                metric_keys.append(key)

    summary: dict[str, Any] = {}
    for key in metric_keys:
        values = np.array([float(metrics.get(key, float("nan"))) for _, metrics in runs], dtype=np.float64)
        finite = values[np.isfinite(values)]
        if finite.size == 0:
            summary[key] = {
                "mean": float("nan"),
                "std": float("nan"),
                "min": float("nan"),
                "max": float("nan"),
                "values": values.tolist(),
                "finite_runs": 0,
            }
            continue
        summary[key] = {
            "mean": float(np.mean(finite)),
            "std": float(np.std(finite, ddof=1)) if finite.size > 1 else 0.0,
            "min": float(np.min(finite)),
            "max": float(np.max(finite)),
            "values": values.tolist(),
            "finite_runs": int(finite.size),
        }

    consistent: dict[str, Any] = {}
    for key in runs[0][1]:
        if not key.startswith("_") and isinstance(runs[0][1][key], str):
            values = {metrics.get(key) for _, metrics in runs}
            consistent[key] = runs[0][1][key]
            if len(values) > 1:
                consistent[f"{key}_INCONSISTENT"] = sorted(str(value) for value in values)

    return {
        "num_runs": len(runs),
        "seeds": [int(seed) for seed, _ in runs],
        "summary": summary,
        "constant": consistent,
        "per_run": [{"seed": int(seed), "metrics": metrics} for seed, metrics in runs],
    }


def build_upper_limit_matrix(
    per_env: dict[str, list[float]],
    level_ids: np.ndarray,
    type_ids: np.ndarray,
    terrain_types: Sequence[str],
    spec: TerrainSpec,
    num_levels: int,
) -> dict[str, Any]:
    """把逐环境结果汇总成网格。

    Args:
        per_env: 含 ``success`` / ``max_progress`` / ``max_tilt`` 的逐环境列表，
            由 :func:`eval_kit.core.metrics.compute_metrics` 在
            ``include_per_env=True`` 时给出。
        level_ids: ``(N,)`` 每个环境所属的难度等级（0 起）。
        type_ids: ``(N,)`` 每个环境所属的地形类型序号（对应 ``terrain_types``）。
        terrain_types: 地形类型名，同时决定矩阵列顺序。
        spec: 地形规格，用于取每级的难度值。
        num_levels: 难度等级数（矩阵行数）。
    """
    success = np.asarray(per_env["success"], dtype=np.float64)
    max_progress = np.asarray(per_env["max_progress"], dtype=np.float64)
    max_tilt = np.asarray(per_env["max_tilt"], dtype=np.float64)
    level_ids = np.asarray(level_ids, dtype=np.int64)
    type_ids = np.asarray(type_ids, dtype=np.int64)

    if not (success.shape == max_progress.shape == max_tilt.shape == level_ids.shape == type_ids.shape):
        raise ValueError("逐环境数组与网格编号数组长度不一致。")

    per_type: dict[str, Any] = {}
    for type_index, terrain_type in enumerate(terrain_types):
        level_values = spec.levels_for(terrain_type, num_levels)
        rates: list[float] = []
        level_metrics: list[dict[str, Any]] = []
        for level_index, level_value in enumerate(level_values):
            mask = (type_ids == type_index) & (level_ids == level_index)
            total = int(np.count_nonzero(mask))
            if total == 0:
                rates.append(float("nan"))
                level_metrics.append(
                    {
                        "level": level_index + 1,
                        "difficulty": float(level_value),
                        "success_rate": float("nan"),
                        "success_count": 0.0,
                        "num_envs": 0,
                        "avg_progress_m": float("nan"),
                        "avg_max_tilt_rad": float("nan"),
                    }
                )
                continue
            count = float(np.count_nonzero(success[mask] > 0.5))
            rates.append(count / total)
            level_metrics.append(
                {
                    "level": level_index + 1,
                    "difficulty": float(level_value),
                    "success_rate": count / total,
                    "success_count": count,
                    "num_envs": total,
                    "avg_progress_m": float(np.mean(max_progress[mask])),
                    "avg_max_tilt_rad": float(np.mean(max_tilt[mask])),
                }
            )

        per_type[terrain_type] = {
            "levels": level_metrics,
            "success_rate": rates,
            "success_rate_accumulated": accumulated_success(rates),
            "score": upper_limit_score(rates),
            "max_level": len(level_values),
            "difficulty_unit": spec.unit.get(terrain_type, ""),
        }

    # 打成便于打印与比较的矩阵：行 = 难度等级，列 = 地形类型
    success_matrix: list[list[float]] = []
    accumulated_matrix: list[list[float]] = []
    difficulty_matrix: list[list[float]] = []
    for level_index in range(num_levels):
        success_row: list[float] = []
        accumulated_row: list[float] = []
        difficulty_row: list[float] = []
        for terrain_type in terrain_types:
            entry = per_type[terrain_type]
            rates = entry["success_rate"]
            accumulated = entry["success_rate_accumulated"]
            levels = entry["levels"]
            success_row.append(rates[level_index] if level_index < len(rates) else float("nan"))
            accumulated_row.append(accumulated[level_index] if level_index < len(accumulated) else float("nan"))
            difficulty_row.append(levels[level_index]["difficulty"] if level_index < len(levels) else float("nan"))
        success_matrix.append(success_row)
        accumulated_matrix.append(accumulated_row)
        difficulty_matrix.append(difficulty_row)

    return {
        "benchmark": spec.name,
        "terrain_types": list(terrain_types),
        "num_levels": num_levels,
        "difficulty_unit": {terrain_type: spec.unit.get(terrain_type, "") for terrain_type in terrain_types},
        "per_type": per_type,
        "scores": {terrain_type: per_type[terrain_type]["score"] for terrain_type in terrain_types},
        "matrix": {
            "success_rate": success_matrix,
            "success_rate_accumulated": accumulated_matrix,
            "difficulty": difficulty_matrix,
        },
    }