# -*- coding: utf-8 -*-
"""
第3周 · 实战3（value-based / Q-Learning）—— 评估 + 录视频脚本
=============================================================
对应 PPT 第 36–37 页（评估方法、Q 表可视化、模型文件）与第 28 页学习目标第 4 条：
  "能够将训练好的智能体和代码上传到 Hub，并附上一个漂亮的视频回放和评估分数"

做三件事：
  1) 加载训练好的模型文件（models/qtable_*.pkl，里面同时存着 Q 表和全套配置）
  2) 用贪婪策略独立评估 N 局，输出评估分数（平均回报 ± 标准差、成功率、平均步数）
     -> eval_results/eval_<run>.json
  3) 回放录像：每一帧都标注"第几局/第几步/动作/奖励/Q 值条形图/累计回报"，
     片头有标题卡，片尾有分数卡
     -> videos/replay_<run>.mp4

用法：
  D:\\study\\robot\\pyenv\\DRL\\python.exe evaluate_and_record.py --model models/qtable_frozenlake4x4.pkl
  D:\\study\\robot\\pyenv\\DRL\\python.exe evaluate_and_record.py --model models/qtable_taxi.pkl --episodes 100
  # 想录得更慢/更多局：
  python evaluate_and_record.py --model models/qtable_taxi.pkl --video-episodes 10 --slowmo 8
"""

import argparse
import json
import os
import pickle
import time
from datetime import datetime

import gymnasium as gym
import imageio
import numpy as np
from PIL import Image, ImageDraw, ImageFont

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
EVAL_DIR = os.path.join(BASE_DIR, "eval_results")
VIDEO_DIR = os.path.join(BASE_DIR, "videos")

# ----------------------------------------------------------------------
# 录像画面参数（1280x720 / 30fps，B 站、Hub、PPT 里都合适）
# ----------------------------------------------------------------------
WIDTH, HEIGHT = 1280, 720
HEADER_H = 78
FOOTER_H = 52
FPS = 30
SLOWMO = 6                 # 每个动作帧重复次数（让 Q-Learning 这种短回合看得清）
HOLD_LAST = 1.2            # 每局结束后的定格秒数
TITLE_SECONDS = 2.5
END_SECONDS = 4.0

# 配色
BG = (15, 23, 42)
PANEL = (30, 41, 59)
PANEL_EDGE = (51, 65, 85)
TEXT = (226, 232, 240)
MUTED = (148, 163, 184)
ACCENT = (56, 189, 248)
GREEN = (34, 197, 94)
RED = (239, 68, 68)
AMBER = (245, 158, 11)

FONT_DIR = r"C:\Windows\Fonts"


def load_font(size, bold=False):
    """优先微软雅黑，找不到再退回黑体/默认字体。"""
    names = ["msyhbd.ttc", "msyh.ttc"] if bold else ["msyh.ttc", "simhei.ttf"]
    for name in names:
        path = os.path.join(FONT_DIR, name)
        if os.path.exists(path):
            try:
                return ImageFont.truetype(path, size)
            except Exception:
                pass
    return ImageFont.load_default()


FONT_TITLE = load_font(40, bold=True)
FONT_H1 = load_font(28, bold=True)
FONT_H2 = load_font(22, bold=True)
FONT_BODY = load_font(20)
FONT_SMALL = load_font(16)
FONT_TINY = load_font(14)

# 动作名称（PPT 第 32 页 / 第 38 页）
ACTION_NAMES = {
    "FrozenLake-v1": ["← 向左 (0)", "↓ 向下 (1)", "→ 向右 (2)", "↑ 向上 (3)"],
    "Taxi-v3": ["↓ 向南 (0)", "↑ 向北 (1)", "→ 向东 (2)", "← 向西 (3)",
                "P 搭载乘客 (4)", "D 送达乘客 (5)"],
}
ACTION_SHORT = {
    "FrozenLake-v1": ["←", "↓", "→", "↑"],
    "Taxi-v3": ["南", "北", "东", "西", "载", "送"],
}

# Taxi-v3 地图上的四个上下客点（R 红 / G 绿 / Y 黄 / B 蓝）
TAXI_LOCS = [(0, 0), (0, 4), (4, 0), (4, 3)]
TAXI_LOC_NAMES = ["R 红", "G 绿", "Y 黄", "B 蓝"]


# ======================================================================
# 模型 / 环境
# ======================================================================
def load_model(path):
    with open(path, "rb") as f:
        model = pickle.load(f)
    if "qtable" not in model:
        raise ValueError(f"{path} 里没有 qtable 字段，不是本脚本保存的模型")
    return model


def make_env(model, render_mode="rgb_array"):
    kwargs = dict(model.get("env_kwargs") or {})
    kwargs.pop("render_mode", None)
    return gym.make(model["env_id"], render_mode=render_mode, **kwargs)


def greedy_action(q_table, state):
    return int(np.argmax(q_table[state]))


# ======================================================================
# 评估：贪婪策略跑 N 局（PPT 第 36 页 evaluate_agent 的同一件事）
# ======================================================================
def evaluate(env, model, n_episodes, seeds):
    q = model["qtable"]
    max_steps = model["max_steps"]
    rows = []
    for i in range(n_episodes):
        if seeds is not None and i < len(seeds):
            state, _ = env.reset(seed=int(seeds[i]))
        else:
            state, _ = env.reset()
        total, steps = 0.0, 0
        for step in range(max_steps):
            action = greedy_action(q, state)
            state, reward, terminated, truncated, _ = env.step(action)
            total += float(reward)
            steps = step + 1
            if terminated or truncated:
                break
        rows.append(dict(episode=i + 1, seed=None if seeds is None else int(seeds[i]),
                         reward=round(total, 4), steps=steps, success=int(total > 0)))
    rewards = np.array([r["reward"] for r in rows], dtype=float)
    success = np.array([r["success"] for r in rows], dtype=float)
    steps = np.array([r["steps"] for r in rows], dtype=float)
    result = dict(
        run=model.get("run_name", "unknown"),
        env_id=model["env_id"],
        env_kwargs=model.get("env_kwargs", {}),
        n_training_episodes=model.get("n_training_episodes"),
        n_episodes=n_episodes,
        mean_reward=float(rewards.mean()),
        std_reward=float(rewards.std()),
        min_reward=float(rewards.min()),
        max_reward=float(rewards.max()),
        success_rate=float(success.mean()),
        mean_steps=float(steps.mean()),
        policy="greedy (argmax Q)",
        model_file=os.path.relpath(model["_path"], BASE_DIR) if "_path" in model else None,
        qtable_shape=list(q.shape),
        qtable_nonzero=int(np.count_nonzero(q)),
        evaluated_at=datetime.now().isoformat(timespec="seconds"),
        episodes=rows,
    )
    return result


# ======================================================================
# 画面工具
# ======================================================================
def rounded_rect(draw, box, radius=16, fill=PANEL, outline=PANEL_EDGE, width=2):
    draw.rounded_rectangle(box, radius=radius, fill=fill, outline=outline, width=width)


def draw_text(draw, xy, text, font=FONT_BODY, fill=TEXT, anchor="la"):
    draw.text(xy, text, font=font, fill=fill, anchor=anchor)


def fit_image(frame, max_w, max_h, nearest=True):
    """把 env.render() 的一帧缩放后居中贴到画布上（像素风用最近邻更清晰）。"""
    img = Image.fromarray(np.asarray(frame, dtype=np.uint8))
    scale = min(max_w / img.width, max_h / img.height)
    new_size = (max(1, int(img.width * scale)), max(1, int(img.height * scale)))
    resample = Image.NEAREST if nearest else Image.LANCZOS
    return img.resize(new_size, resample)


def base_canvas():
    canvas = Image.new("RGB", (WIDTH, HEIGHT), BG)
    draw = ImageDraw.Draw(canvas)
    # 顶栏
    draw.rectangle([0, 0, WIDTH, HEADER_H], fill=(15, 23, 42))
    draw.line([0, HEADER_H, WIDTH, HEADER_H], fill=PANEL_EDGE, width=2)
    # 底栏
    draw.rectangle([0, HEIGHT - FOOTER_H, WIDTH, HEIGHT], fill=(15, 23, 42))
    draw.line([0, HEIGHT - FOOTER_H, WIDTH, HEIGHT - FOOTER_H], fill=PANEL_EDGE, width=2)
    return canvas, draw


def title_card(model, result):
    """片头标题卡。"""
    canvas, draw = base_canvas()
    draw_text(draw, (WIDTH // 2, 150), "Q-Learning 智能体 · 训练后回放", FONT_TITLE,
              fill=ACCENT, anchor="mm")
    draw_text(draw, (WIDTH // 2, 215), model.get("title", model["env_id"]), FONT_H1,
              fill=TEXT, anchor="mm")
    draw.line([WIDTH // 2 - 220, 260, WIDTH // 2 + 220, 260], fill=PANEL_EDGE, width=2)

    # 中间：三个信息块
    blocks = [
        ("算法", "Q-Learning（从零手写 Q 表）"),
        ("训练规模", f"{model['n_training_episodes']:,} 回合"),
        ("Q 表", f"{model['qtable'].shape[0]} 状态 × {model['qtable'].shape[1]} 动作"),
        ("超参数", f"lr={model['learning_rate']}  γ={model['gamma']}  "
                   f"ε {model['max_epsilon']}→{model['min_epsilon']}"),
    ]
    y = 310
    for k, v in blocks:
        draw_text(draw, (WIDTH // 2 - 300, y), f"{k}", FONT_SMALL, fill=MUTED)
        draw_text(draw, (WIDTH // 2 - 160, y), f"{v}", FONT_BODY, fill=TEXT)
        y += 46

    draw_text(draw, (WIDTH // 2, 570),
              f"独立评估 {result['n_episodes']} 局："
              f"平均回报 {result['mean_reward']:.2f} ± {result['std_reward']:.2f}   "
              f"成功率 {result['success_rate'] * 100:.0f}%",
              FONT_H2, fill=GREEN, anchor="mm")
    draw_text(draw, (WIDTH // 2, 615),
              "对应《第3周-深度强化学习原理实践》PPT 第 27–40 页 · value-based 实践",
              FONT_SMALL, fill=MUTED, anchor="mm")
    draw_text(draw, (WIDTH // 2, 640), "模型文件：models/" + os.path.basename(model["_path"]),
              FONT_SMALL, fill=MUTED, anchor="mm")
    return canvas


def score_card(model, result):
    """片尾分数卡。"""
    canvas, draw = base_canvas()
    draw_text(draw, (WIDTH // 2, 120), "评估结果", FONT_TITLE, fill=ACCENT, anchor="mm")
    draw_text(draw, (WIDTH // 2, 185), model.get("title", model["env_id"]), FONT_H1,
              fill=TEXT, anchor="mm")

    # 大号分数
    mean, std = result["mean_reward"], result["std_reward"]
    draw_text(draw, (WIDTH // 2 - 200, 300), "平均回报", FONT_BODY, fill=MUTED, anchor="mm")
    draw_text(draw, (WIDTH // 2 - 200, 360), f"{mean:.2f}", load_font(64, bold=True),
              fill=GREEN if mean > 0 else RED, anchor="mm")
    draw_text(draw, (WIDTH // 2 - 200, 410), f"± {std:.2f}", FONT_H2, fill=MUTED, anchor="mm")

    draw_text(draw, (WIDTH // 2 + 60, 300), "成功率", FONT_BODY, fill=MUTED, anchor="mm")
    draw_text(draw, (WIDTH // 2 + 60, 360), f"{result['success_rate'] * 100:.0f}%",
              load_font(64, bold=True), fill=ACCENT, anchor="mm")

    draw_text(draw, (WIDTH // 2 + 300, 300), "平均步数", FONT_BODY, fill=MUTED, anchor="mm")
    draw_text(draw, (WIDTH // 2 + 300, 360), f"{result['mean_steps']:.1f}",
              load_font(64, bold=True), fill=AMBER, anchor="mm")

    draw_text(draw, (WIDTH // 2, 505),
              f"{result['n_episodes']} 局独立评估 · 贪婪策略 argmax_a Q(s,a)",
              FONT_BODY, fill=TEXT, anchor="mm")
    # 迷你分数分布
    rewards = [e["reward"] for e in result["episodes"]]
    values, counts = np.unique(np.round(rewards, 2), return_counts=True)
    draw_text(draw, (WIDTH // 2, 555),
              "逐局回报分布：" + "  ".join(f"{v:g}×{c}" for v, c in zip(values[:8], counts[:8]))
              + (" …" if len(values) > 8 else ""),
              FONT_SMALL, fill=MUTED, anchor="mm")
    draw_text(draw, (WIDTH // 2, 620),
              "模型与代码已打包上传 Hub：附视频回放 + 评估分数",
              FONT_SMALL, fill=MUTED, anchor="mm")
    return canvas


def draw_q_bars(draw, q_row, action_names, chosen_action, x, y, width):
    """把当前状态各动作的 Q 值画成水平条形图（直观看出贪心策略选了什么）。"""
    q = np.asarray(q_row, dtype=float)
    lo, hi = float(q.min()), float(q.max())
    span = hi - lo if hi - lo > 1e-9 else 1.0
    bar_h = 26
    gap = 10
    for a in range(len(q)):
        yy = y + a * (bar_h + gap)
        chosen = a == chosen_action
        fill = ACCENT if chosen else (71, 85, 105)
        # 动作标签
        draw_text(draw, (x, yy + bar_h // 2), action_names[a], FONT_SMALL,
                  fill=TEXT if chosen else MUTED, anchor="lm")
        # 条形
        left = x + 150
        draw.rectangle([left, yy, left + width, yy + bar_h], fill=(51, 65, 85))
        frac = (q[a] - lo) / span
        draw.rectangle([left, yy, left + max(3, int(width * frac)), yy + bar_h], fill=fill)
        # 数值
        draw_text(draw, (left + width + 12, yy + bar_h // 2), f"{q[a]:.3f}", FONT_SMALL,
                  fill=ACCENT if chosen else MUTED, anchor="lm")
        if chosen:
            draw_text(draw, (left + width + 78, yy + bar_h // 2), "← 最优", FONT_SMALL,
                      fill=ACCENT, anchor="lm")


def game_frame(model, env, state, step, episode_idx, n_video_eps, total_reward,
               last_reward, action, ended, success, eval_result):
    """合成一帧：左边是环境画面，右边是 Q 值/动作/奖励信息。"""
    canvas, draw = base_canvas()

    title = model.get("title", model["env_id"])
    draw_text(draw, (32, HEADER_H // 2), f"Q-Learning 智能体回放 · {title}",
              FONT_H2, fill=TEXT, anchor="lm")
    draw_text(draw, (WIDTH - 32, HEADER_H // 2),
              "PPT 第 27–40 页 · value-based", FONT_SMALL, fill=MUTED, anchor="rm")

    # ---------- 左边：环境画面 ----------
    box = [32, HEADER_H + 22, 672, HEIGHT - FOOTER_H - 22]
    rounded_rect(draw, box)
    raw = env.render()
    if raw is not None:
        inner_w = box[2] - box[0] - 24
        inner_h = box[3] - box[1] - 24
        img = fit_image(raw, inner_w, inner_h)
        px = box[0] + (box[2] - box[0] - img.width) // 2
        py = box[1] + (box[3] - box[1] - img.height) // 2
        canvas.paste(img, (px, py))

    # ---------- 右边：信息面板 ----------
    px0, px1 = 700, WIDTH - 32
    rounded_rect(draw, [px0, HEADER_H + 22, px1, HEIGHT - FOOTER_H - 22])

    x = px0 + 24
    y = HEADER_H + 50
    draw_text(draw, (x, y), f"第 {episode_idx + 1} / {n_video_eps} 局", FONT_H2, fill=TEXT)
    y += 34

    # Taxi 额外显示乘客/目的地
    if model["env_id"] == "Taxi-v3":
        taxi_row, taxi_col, pass_idx, dest_idx = env.unwrapped.decode(int(state))
        pass_text = "在车上" if pass_idx == 4 else TAXI_LOC_NAMES[pass_idx] + " 处"
        draw_text(draw, (x, y), f"出租车  ({taxi_row},{taxi_col})  乘客 {pass_text}  "
                                f"目的地 {TAXI_LOC_NAMES[dest_idx]}", FONT_SMALL, fill=MUTED)
        y += 30
    else:
        draw_text(draw, (x, y), f"状态编号 s = {int(state)}", FONT_SMALL, fill=MUTED)
        y += 30

    # 步数 / 动作 / 奖励
    draw_text(draw, (x, y), "步数", FONT_SMALL, fill=MUTED)
    draw_text(draw, (x + 70, y), f"{step}", FONT_H2, fill=TEXT)
    draw_text(draw, (x + 180, y), "动作", FONT_SMALL, fill=MUTED)
    draw_text(draw, (x + 250, y),
              ACTION_NAMES[model["env_id"]][action] if action is not None else "—",
              FONT_H2, fill=ACCENT)
    y += 38

    draw_text(draw, (x, y), "单步奖励", FONT_SMALL, fill=MUTED)
    draw_text(draw, (x + 120, y), f"{last_reward:+.2f}" if last_reward is not None else "—",
              FONT_H2, fill=GREEN if (last_reward or 0) > 0 else TEXT)
    draw_text(draw, (x + 260, y), "累计回报", FONT_SMALL, fill=MUTED)
    draw_text(draw, (x + 360, y), f"{total_reward:+.2f}", FONT_H2,
              fill=GREEN if total_reward > 0 else TEXT)
    y += 44

    # Q 值条形图（终局状态不再更新 Q，改成提示文字）
    if ended:
        draw_text(draw, (x, y), "终局状态：回合结束，Q 表在这一帧不再更新", FONT_SMALL, fill=MUTED)
    else:
        draw_text(draw, (x, y), f"Q(s={int(state)}, ·)  ← 贪心策略从这里选动作", FONT_SMALL, fill=MUTED)
        y += 30
        q = model["qtable"][int(state)]
        chosen = int(np.argmax(q))
        bars_width = 180
        draw_q_bars(draw, q, ACTION_SHORT[model["env_id"]], chosen, x, y, bars_width)

    # 回合结束横幅
    if ended:
        banner_y = HEIGHT - FOOTER_H - 105
        if model["env_id"] == "Taxi-v3":
            text, color = ("√ 成功送达乘客（+20 奖励）", GREEN) if success else \
                          ("× 超时 / 未送达（回合结束）", RED)
        else:
            text, color = ("√ 成功到达目标（+1 奖励）", GREEN) if success else \
                          ("× 掉进冰洞（奖励 0，回合结束）", RED)
        draw.rounded_rectangle([px0 + 20, banner_y, px1 - 20, banner_y + 62],
                               radius=12, fill=(color[0] // 4, color[1] // 4, color[2] // 4),
                               outline=color, width=2)
        draw_text(draw, ((px0 + px1) // 2, banner_y + 31), text, FONT_H2, fill=color, anchor="mm")

    # ---------- 底栏：总体评估分数 ----------
    draw_text(draw, (32, HEIGHT - FOOTER_H // 2),
              f"评估 {eval_result['n_episodes']} 局：平均回报 "
              f"{eval_result['mean_reward']:.2f} ± {eval_result['std_reward']:.2f}   "
              f"成功率 {eval_result['success_rate'] * 100:.0f}%   "
              f"平均步数 {eval_result['mean_steps']:.1f}",
              FONT_SMALL, fill=MUTED, anchor="lm")
    draw_text(draw, (WIDTH - 32, HEIGHT - FOOTER_H // 2), "贪心策略 argmax_a Q(s,a)",
              FONT_SMALL, fill=MUTED, anchor="rm")
    return canvas


# ======================================================================
# 录像
# ======================================================================
def record_video(env, model, result, out_path, n_video_episodes, slowmo, fps):
    q = model["qtable"]
    max_steps = model["max_steps"]
    seeds = model.get("eval_seed")
    writer = imageio.get_writer(out_path, fps=fps, codec="libx264", quality=8,
                                macro_block_size=2)
    n_frames = 0

    def push(img, repeat=1):
        nonlocal n_frames
        frame = np.asarray(img, dtype=np.uint8)
        for _ in range(repeat):
            writer.append_data(frame)
            n_frames += 1

    # 片头
    push(title_card(model, result), int(TITLE_SECONDS * fps))
    hold_end = int(HOLD_LAST * fps)

    for ep in range(n_video_episodes):
        if seeds is not None and ep < len(seeds):
            state, _ = env.reset(seed=int(seeds[ep]))
        else:
            state, _ = env.reset()

        total_reward = 0.0
        last_reward = None
        prev_reward = None
        last_action = None
        step = 0
        for step in range(max_steps):
            action = greedy_action(q, state)
            last_action = action
            new_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            total_reward += float(reward)
            success = total_reward > 0
            last_reward = float(reward)
            # 这一步的"决策画面"：画面里的状态是"还没执行这个动作之前"的状态，
            # 所以累计回报用执行前的值，单步奖励显示上一步拿到的奖励
            push(game_frame(model, env, state, step, ep, n_video_episodes,
                            total_reward - last_reward, prev_reward, action, False, success, result),
                 slowmo)
            prev_reward = last_reward
            state = new_state
            if done:
                break
        # 回合结束：定格 + 结束横幅
        ended_frame = game_frame(model, env, state, step + 1, ep, n_video_episodes,
                                 total_reward, last_reward, last_action, True,
                                 total_reward > 0, result)
        push(ended_frame, hold_end)
        print(f"  [录像] 第 {ep + 1}/{n_video_episodes} 局：回报={total_reward:+.2f}，"
              f"步数={step + 1}，{'√ 成功' if total_reward > 0 else '× 失败'}")

    # 片尾分数卡
    push(score_card(model, result), int(END_SECONDS * fps))
    writer.close()
    return n_frames


# ======================================================================
# 主流程
# ======================================================================
def main():
    parser = argparse.ArgumentParser(description="评估 Q-Learning 智能体并录制回放视频")
    parser.add_argument("--model", type=str, required=True, help="models/qtable_*.pkl")
    parser.add_argument("--episodes", type=int, default=None,
                        help="评估局数（默认用模型里记录的 n_eval_episodes）")
    parser.add_argument("--video-episodes", type=int, default=8, help="录几局回放（默认 8）")
    parser.add_argument("--slowmo", type=int, default=SLOWMO, help="每个动作帧重复几次（默认 6）")
    parser.add_argument("--fps", type=int, default=FPS, help="视频帧率（默认 30）")
    parser.add_argument("--out-video", type=str, default=None, help="视频输出路径（.mp4）")
    parser.add_argument("--out-json", type=str, default=None, help="评估结果输出路径（.json）")
    parser.add_argument("--skip-video", action="store_true",
                        help="只评估、不录视频（改完分数想快速重算时用）")
    args = parser.parse_args()

    model = load_model(args.model)
    model["_path"] = os.path.abspath(args.model)
    model["run_name"] = os.path.splitext(os.path.basename(args.model))[0].replace("qtable_", "")
    if not model.get("title"):
        model["title"] = model["env_id"]

    os.makedirs(EVAL_DIR, exist_ok=True)
    os.makedirs(VIDEO_DIR, exist_ok=True)

    run = model["run_name"]
    out_json = args.out_json or os.path.join(EVAL_DIR, f"eval_{run}.json")
    out_video = args.out_video or os.path.join(VIDEO_DIR, f"replay_{run}.mp4")

    n_eval = args.episodes or model.get("n_eval_episodes", 100)
    seeds = model.get("eval_seed")
    if seeds is not None and len(seeds) < n_eval:
        seeds = None   # 种子不够就不指定，退回随机起点

    print("=" * 78)
    print(f"评估 + 录像 · {model['title']}")
    print(f"模型文件：{os.path.relpath(model['_path'], BASE_DIR)}")
    print(f"贪婪策略：argmax_a Q(s,a)   Q 表形状：{model['qtable'].shape}")
    print("=" * 78)

    env = make_env(model)
    print(f"\n[评估] 用贪婪策略跑 {n_eval} 局（独立环境，起点种子："
          f"{'模型自带的 ' + str(len(seeds)) + ' 个种子' if seeds is not None else '随机'}）")
    started = time.time()
    result = evaluate(env, model, n_eval, seeds)
    result["evaluation_seconds"] = round(time.time() - started, 2)
    print(f"  Mean_reward = {result['mean_reward']:.2f} +/- {result['std_reward']:.2f}")
    print(f"  成功率      = {result['success_rate'] * 100:.0f}%")
    print(f"  平均步数    = {result['mean_steps']:.1f}")
    print(f"  回报范围    = [{result['min_reward']:.2f}, {result['max_reward']:.2f}]")
    with open(out_json, "w", encoding="utf-8") as f:
        json.dump(result, f, ensure_ascii=False, indent=2)
    print(f"[输出] 评估分数 = {os.path.relpath(out_json, BASE_DIR)}")

    if args.skip_video:
        print("\n[跳过录像]（--skip-video）")
        env.close()
        print("\n" + "=" * 78)
        print(f"完成。评估分数 {result['mean_reward']:.2f} ± {result['std_reward']:.2f}，"
              f"成功率 {result['success_rate'] * 100:.0f}%")
        print("=" * 78)
        return

    print(f"\n[录像] 录 {args.video_episodes} 局回放，每帧慢放 ×{args.slowmo}，{args.fps}fps")
    n_frames = record_video(env, model, result, out_video, args.video_episodes,
                            args.slowmo, args.fps)
    secs = n_frames / args.fps
    print(f"[输出] 回放视频 = {os.path.relpath(out_video, BASE_DIR)} "
          f"({n_frames} 帧 / {secs:.1f} 秒 / {os.path.getsize(out_video) / 1024:.0f} KB)")
    env.close()

    print("\n" + "=" * 78)
    print(f"完成。评估分数 {result['mean_reward']:.2f} ± {result['std_reward']:.2f}，"
          f"成功率 {result['success_rate'] * 100:.0f}%")
    print("=" * 78)


if __name__ == "__main__":
    main()