"""用相同的一组 DelivCoal 场景评估多个 MARL-MRT2A checkpoint。"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parents[2]
OFFICIAL_SRC = REPO_ROOT / "third_party/marl_mrt2a/marl_mrt2a/marl/src"
DEFAULT_ENV_CONFIG = REPO_ROOT / "experiments/reproduction/delivcoal_formal_base.yaml"
sys.path.insert(0, str(OFFICIAL_SRC))

import main as epymarl_main  # noqa: E402


def scalar(value):
    return float(value.item()) if hasattr(value, "item") else float(value)


def evaluate_one(model_root, checkpoint, env_config, eval_seed,
                 eval_episodes, batch_size, episode_limit, grid_size, cuda,
                 observation_layout="legacy20", share_intention="none"):
    """加载一个训练步对应的模型，并运行指定数量的评估 episode。"""
    overrides = {
        "env_args.key": repr("delivcoal-formal"),
        "env_args.config_path": repr(str(env_config)),
        "env_args.grid_size": grid_size,
        "env_args.time_limit": episode_limit,
        "env_args.observation_layout": repr(observation_layout),
        "env_args.share_intention": repr(share_intention),
        "critic_type": repr("cnn_cv_critic"),
        "critic_arch": repr("unet,8,1,2&batchNorm1d;linear,50;relu"),
        "agent": repr("cnn"),
        "agent_arch": repr("unet,8,1,2&"),
        "action_grid": True,
        "agent_distance_exp": 1.0,
        "current_target_factor": 1.0,
        "agent_reeval_rate": True,
        "filter_avail_by_objects": False,
        "gamma": 0.95,
        "standardise_returns": False,
        "standardise_rewards": True,
        "q_nstep": 5,
        "obs_agent_id": False,
        "obs_last_action": False,
        "batch_size_run": batch_size,
        "test_nepisode": eval_episodes,
        "checkpoint_path": repr(str(model_root)),
        "load_step": checkpoint,
        "evaluate": True,
        "use_cuda": cuda,
        "seed": eval_seed,
        "name": repr(f"marl_mrt2a_delivcoal_eval_c{checkpoint}_s{eval_seed}"),
        "wandb_sweep": False,
        "use_tensorboard": False,
    }
    argv = ["main.py", "--config=mappo", "--env-config=delivcoal", "with"]
    argv.extend(f"{key}={value}" for key, value in overrides.items())
    runner, _, _, _, logger = epymarl_main.main_from_arg(argv, just_build=True)
    # runner 每次同时运行 batch_size 个环境，因此向上取整调用次数。
    runs = max(1, (eval_episodes + runner.batch_size - 1) // runner.batch_size)
    for _ in range(runs):
        runner.run(test_mode=True)
    runner.close_env()
    return {
        key: scalar(values[-1][1])
        for key, values in logger.stats.items()
        if key.startswith("test_") and values
    } | {
        "loaded_step": int(runner.t_env),
        "evaluated_episodes": int(runs * runner.batch_size),
    }


def run(args):
    """依次评估 checkpoint，并把同口径结果写入一个 JSON 文件。"""
    run_dir = Path(args.run_dir).resolve()
    model_roots = list((run_dir / "results/models").glob("*"))
    if len(model_roots) != 1:
        raise RuntimeError(f"Expected one model root, found {len(model_roots)}")
    output = Path(args.output).resolve() if args.output else run_dir / "heldout_eval"
    output.mkdir(parents=True, exist_ok=True)
    env_config = Path(args.env_config).resolve()
    results = {}
    for checkpoint in args.checkpoints:
        result = evaluate_one(
            model_roots[0], checkpoint, env_config, args.eval_seed,
            args.eval_episodes, args.batch_size_run, args.episode_limit,
            args.grid_size, args.cuda,
            args.observation_layout, args.share_intention,
        )
        results[str(checkpoint)] = result
        print(
            f"MARL-MRT2A DelivCoal {checkpoint}: "
            f"success={result.get('test_success_mean')}, "
            f"completion={result.get('test_task_completion_ratio_mean')}, "
            f"sent={result.get('test_sent_percentage_mean')}",
            flush=True,
        )
    payload = {
        "method": "MARL-MRT2A",
        "environment": "DelivCoal",
        "protocol": {
            "eval_seed": args.eval_seed,
            "eval_episodes": args.eval_episodes,
            "batch_size_run": args.batch_size_run,
            "paired_base_seed": True,
            "observation_layout": args.observation_layout,
            "share_intention": args.share_intention,
        },
        "checkpoints": results,
    }
    (output / "summary.json").write_text(
        json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8"
    )


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--run-dir", required=True)
    parser.add_argument("--output")
    parser.add_argument("--env-config", default=str(DEFAULT_ENV_CONFIG))
    parser.add_argument("--checkpoints", nargs="+", type=int, required=True)
    parser.add_argument("--eval-seed", type=int, default=40_000)
    parser.add_argument("--eval-episodes", type=int, default=128)
    parser.add_argument("--batch-size-run", type=int, default=2)
    parser.add_argument("--episode-limit", type=int, default=750)
    parser.add_argument("--grid-size", type=int, default=21)
    parser.add_argument(
        "--observation-layout", choices=("legacy20", "full27"),
        default="legacy20",
    )
    parser.add_argument(
        "--share-intention", choices=("none", "path"), default="none"
    )
    parser.add_argument("--cuda", action="store_true")
    return parser.parse_args()


if __name__ == "__main__":
    run(parse_args())