"""一键异步启动全部训练。"""
import os
import subprocess
import sys
from datetime import datetime

AGENTS = ["ppo", "appo", "mappo", "dqn", "ddpg", "td3"]


def main():
    run_timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    processes = {}

    for agent in AGENTS:
        print(f"[{datetime.now().strftime('%H:%M:%S')}] 启动 {agent.upper()} 训练...")
        log_dir = os.path.join("logs", "multi-model", run_timestamp, agent)
        checkpoint_dir = os.path.join("checkpoints", "multi-model", run_timestamp, agent)
        os.makedirs(log_dir, exist_ok=True)
        os.makedirs(checkpoint_dir, exist_ok=True)
        stdout_path = os.path.join(log_dir, "stdout.txt")

        process = subprocess.Popen(
            [
                sys.executable,
                "-m",
                f"training.train_{agent}",
                "--log-dir",
                log_dir,
                "--checkpoint-dir",
                checkpoint_dir,
                "--run-timestamp",
                run_timestamp,
            ],
            stdout=open(stdout_path, "w", encoding="utf-8"),
            stderr=subprocess.STDOUT,
        )
        processes[agent] = process
        print(f"  PID: {process.pid}")

    print(f"\n全部 {len(AGENTS)} 个训练已启动，等待完成...")
    print(f"本次多模型时间戳: {run_timestamp}\n")

    for agent, process in processes.items():
        process.wait()
        status = "完成" if process.returncode == 0 else f"失败(code={process.returncode})"
        print(f"[{agent.upper()}] {status}")


if __name__ == "__main__":
    main()