"""使用 MARL-MRT2A 原项目的 PyMARL 训练流程运行 DelivCoal。"""

from __future__ import annotations

import argparse
import json
import os
import sys
from datetime import datetime
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 run(args):
    """把命令行参数转换成 PyMARL 配置覆盖项并启动训练。"""
    output = Path(args.output).resolve()
    output.mkdir(parents=True, exist_ok=True)
    env_config = Path(args.env_config).resolve()
    name = f"marl_mrt2a_delivcoal_s{args.seed}"
    overrides = {
        # 环境适配器把 DelivCoal 的任务位置投到二维动作网格。
        "env_args.key": repr("delivcoal-formal"),
        "env_args.config_path": repr(str(env_config)),
        "env_args.grid_size": args.grid_size,
        "env_args.time_limit": args.episode_limit,
        "env_args.observation_layout": repr(args.observation_layout),
        "env_args.share_intention": repr(args.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,
        "lr": args.lr,
        "entropy_coef": args.entropy_coef,
        "standardise_returns": False,
        "standardise_rewards": True,
        "q_nstep": 5,
        "obs_agent_id": False,
        "obs_last_action": False,
        "buffer_size": max(args.batch_size, args.buffer_size),
        "batch_size": args.batch_size,
        "batch_size_run": args.batch_size_run,
        "target_update_interval_or_tau": 200,
        "t_max": args.t_max,
        "test_nepisode": args.test_nepisode,
        "test_interval": args.test_interval,
        "log_interval": args.log_interval,
        "runner_log_interval": args.log_interval,
        "learner_log_interval": args.log_interval,
        "save_model": True,
        "save_model_interval": args.save_model_interval,
        "run_perm_importance": False,
        "save_path": repr(str(output)),
        "use_cuda": bool(args.cuda),
        "seed": args.seed,
        "name": repr(name),
        "wandb_sweep": False,
        "use_tensorboard": False,
    }
    argv = ["main.py", "--config=mappo", "--env-config=delivcoal", "with"]
    # PyMARL 使用 Sacred 的 key=value 形式接收运行时覆盖项。
    argv.extend(f"{key}={value}" for key, value in overrides.items())
    metadata = {
        "started_at": datetime.now().isoformat(timespec="seconds"),
        "method": "official MARL-MRT2A CNN policy and target revision",
        "adapter": "DelivCoal spatial task-cell adapter",
        "official_source": str(OFFICIAL_SRC),
        "environment_config": str(env_config),
        "argv": argv,
        "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
    }
    config_path = output / "reproduction_config.json"
    config_path.write_text(
        json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    epymarl_main.main_from_arg(argv)
    metadata["finished_at"] = datetime.now().isoformat(timespec="seconds")
    config_path.write_text(
        json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8"
    )


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--seed", type=int, required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--env-config", default=str(DEFAULT_ENV_CONFIG))
    parser.add_argument("--t-max", type=int, default=3_600_000)
    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("--batch-size-run", type=int, default=4)
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--buffer-size", type=int, default=16)
    parser.add_argument("--test-nepisode", type=int, default=32)
    parser.add_argument("--test-interval", type=int, default=50_000)
    parser.add_argument("--save-model-interval", type=int, default=50_000)
    parser.add_argument("--log-interval", type=int, default=10_000)
    parser.add_argument("--lr", type=float, default=5e-4)
    parser.add_argument("--entropy-coef", type=float, default=0.01)
    parser.add_argument("--cuda", action="store_true")
    return parser.parse_args()


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