"""CFM 和 MARL-MRT2A 复现程序共用的 DelivCoal 工具。"""

from __future__ import annotations

import json
import os
import random
from pathlib import Path
from types import SimpleNamespace

import numpy as np
import torch
import yaml


REPO_ROOT = Path(__file__).resolve().parents[2]
DEFAULT_CONFIG = Path(__file__).with_name("delivcoal_formal_base.yaml")


def load_env_config(path=DEFAULT_CONFIG, **overrides):
    """读取环境配置，并把奖励文件路径转换成可直接使用的绝对路径。"""
    path = Path(path).expanduser().resolve()
    payload = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
    payload.update(overrides)
    rewards = Path(payload.get("rewards_config_file", "rewards.yaml"))
    if not rewards.is_absolute():
        rewards = REPO_ROOT / rewards
    payload["rewards_config_file"] = str(rewards)
    return SimpleNamespace(**payload)


def config_dict(config):
    return {
        key: str(value) if isinstance(value, Path) else value
        for key, value in vars(config).items()
    }


def seed_everything(seed):
    """给 Python、NumPy 和 PyTorch 设置同一个随机种子。"""
    seed = int(seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def atomic_json(path, payload):
    """通过临时文件更新 JSON，避免其他进程读到半份内容。"""
    path = Path(path)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(
        json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    os.replace(temporary, path)


def summarize_episode(info, episode_reward, episode_length, sent_total):
    """把 DelivCoal 的 episode 结果整理成两种复现方法共用的指标。"""
    tasks = info.get("tasks", [])
    # DelivCoal 中类型 1 表示配送任务。这里直接使用数值，避免导入完整环境模块。
    delivery = [task for task in tasks if int(task.get("type", -1)) == 1]
    completed = sum(bool(task.get("completed")) for task in delivery)
    discovered = sum(bool(task.get("discovered")) for task in delivery)
    count = max(1, len(delivery))
    sent_goal = max(1e-8, float(info.get("sent_goal", 0.0)))
    return {
        "success": float(completed == len(delivery) and len(delivery) > 0),
        "task_completion_ratio": float(completed / count),
        "discovered_ratio": float(discovered / count),
        "sent_percentage": float(sent_total / sent_goal),
        "episode_reward": float(episode_reward),
        "episode_length": int(episode_length),
        "coalition_changes": float(
            info.get("reconfiguration", {}).get("total_coalition_changes", 0)
        ),
        "reconfiguration_decisions": float(
            info.get("reconfiguration", {}).get("total_decisions", 0)
        ),
    }


def aggregate_rows(rows):
    """计算多次评估的均值、标准差和 95% 置信区间。"""
    if not rows:
        return {}
    result = {}
    for key in rows[0]:
        values = np.asarray([row[key] for row in rows], dtype=np.float64)
        result[key] = {
            "mean": float(values.mean()),
            "std": float(values.std()),
            "ci95": float(1.96 * values.std(ddof=1) / np.sqrt(len(values)))
            if len(values) > 1
            else 0.0,
        }
    return result