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

from __future__ import annotations

import argparse
import csv
import json
import sys
from pathlib import Path

import torch

REPO_ROOT = Path(__file__).resolve().parents[2]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from experiments.reproduction.delivcoal_adapter_common import (
    DEFAULT_CONFIG,
    aggregate_rows,
    atomic_json,
    load_env_config,
)
from experiments.reproduction.train_cfm_delivcoal import AttentionNet, rollout


def run(args):
    """依次加载 checkpoint，运行评估并汇总为 CSV 和 JSON。"""
    run_dir = Path(args.run_dir).resolve()
    output = Path(args.output).resolve() if args.output else run_dir / "heldout_eval"
    output.mkdir(parents=True, exist_ok=True)
    config = load_env_config(args.env_config)
    device = torch.device(
        f"cuda:{args.cuda}" if args.cuda >= 0 and torch.cuda.is_available() else "cpu"
    )
    seeds = list(range(args.eval_seed, args.eval_seed + args.eval_episodes))
    # seeds 对所有 checkpoint 保持一致。
    summaries = {}
    all_rows = []
    for episode in args.episodes:
        # 每个 checkpoint 使用同一组 seeds，便于直接比较训练阶段差异。
        checkpoint_path = run_dir / "models" / f"actor_ep{episode:06d}.pt"
        checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
        network = AttentionNet(
            int(checkpoint["agent_input_dim"]),
            int(checkpoint["task_input_dim"]),
            int(checkpoint["embedding_dim"]),
        ).to(device)
        network.load_state_dict(checkpoint["model"])
        network.eval()
        # 每个 seed 只创建一个场景，greedy 模式不再引入动作采样随机性。
        rows = []
        with torch.no_grad():
            for seed in seeds:
                row, _ = rollout(
                    network,
                    config,
                    seed,
                    device,
                    not args.sample_policy,
                    policy_seed=(seed + args.policy_seed_offset)
                    if args.sample_policy else None,
                )
                row["checkpoint_episode"] = int(episode)
                rows.append(row)
                all_rows.append(row)
        summaries[str(episode)] = aggregate_rows(rows)
        summary = summaries[str(episode)]
        print(
            f"CFM DelivCoal ep{episode}: success={summary['success']['mean']:.3f}, "
            f"completion={summary['task_completion_ratio']['mean']:.3f}, "
            f"sent={summary['sent_percentage']['mean']:.3f}",
            flush=True,
        )
    with (output / "per_episode.csv").open("w", newline="", encoding="utf-8") as file:
        # CSV 保存每个场景的原始结果，JSON 保存各指标的汇总结果。
        writer = csv.DictWriter(file, fieldnames=list(all_rows[0]))
        writer.writeheader()
        writer.writerows(all_rows)
    atomic_json(output / "summary.json", {
        "method": "CFM",
        "environment": "DelivCoal",
        "protocol": {
            "eval_seed": args.eval_seed,
            "eval_episodes": args.eval_episodes,
            "paired_seeds": True,
            "greedy": not args.sample_policy,
            "policy_seed_offset": (
                args.policy_seed_offset if args.sample_policy else None
            ),
        },
        "checkpoints": summaries,
    })


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_CONFIG))
    parser.add_argument("--episodes", nargs="+", type=int, required=True)
    parser.add_argument("--eval-episodes", type=int, default=128)
    parser.add_argument("--eval-seed", type=int, default=40_000)
    parser.add_argument("--cuda", type=int, default=-1)
    parser.add_argument(
        "--sample-policy",
        action="store_true",
        help="Sample actions reproducibly instead of greedy decoding.",
    )
    parser.add_argument("--policy-seed-offset", type=int, default=1_000_000)
    return parser.parse_args()


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