# -*- coding: utf-8 -*-
"""
第3周 · 实战3（value-based / Q-Learning）—— 训练脚本
=====================================================
对应 PPT《第3周-深度强化学习原理实践》第 27–40 页：
  · 第 29 页：Q 表 / Q 函数 / 最优策略 pi*(s) = argmax_a Q*(s,a)
  · 第 31–32 页：FrozenLake-v1（4x4 不滑版）观测/动作/奖励
  · 第 33 页：Q 表初始化 + 贪心策略
  · 第 34 页：epsilon-greedy 行为策略 + 超参数
  · 第 35 页：Q-Learning 训练循环（Bellman 更新）
  · 第 36–37 页：训练、Q 表可视化、评估
  · 第 38–40 页：Taxi-v3

本脚本从零手写 Q-Learning（不依赖 stable-baselines3），一次跑通三个环境：
  frozenlake4x4  FrozenLake-v1 4x4，is_slippery=False（PPT 第一部分，确定性）
  frozenlake8x8  FrozenLake-v1 8x8，is_slippery=True （PPT 第 31 页提到的进阶版本，随机性）
  taxi           Taxi-v3（PPT 第二部分，500 个状态 × 6 个动作）

输出（全部落在本目录下）：
  logs/<run>/train_console.log   完整控制台日志（含环境自检/超参/每 N 回合进度/每次评估）
  logs/<run>/progress.csv        结构化进度数据（画曲线 / 写报告用）
  logs/<run>/train_curve.png     训练曲线（回报、epsilon、成功率、评估）
  models/qtable_<run>.pkl        训练出的模型文件（Q 表 + 全套配置，可被 evaluate_and_record.py 直接加载）

用法（本机已装好 DRL 环境，无需再装包）：
  D:\\study\\robot\\pyenv\\DRL\\python.exe train_qlearning.py --env frozenlake4x4
  D:\\study\\robot\\pyenv\\DRL\\python.exe train_qlearning.py --env frozenlake8x8
  D:\\study\\robot\\pyenv\\DRL\\python.exe train_qlearning.py --env taxi

只改一个变量做对照实验（例如演示"epsilon 衰减太慢"）：
  python train_qlearning.py --env frozenlake4x4 --decay-rate 0.0001 --run-name fl4x4_slow_decay
"""

import argparse
import csv
import os
import pickle
import random
import sys
import time
from datetime import datetime

import gymnasium as gym
import matplotlib
matplotlib.use("Agg")               # 不弹窗，直接存 png
import matplotlib.pyplot as plt
import numpy as np

# ======================================================================
# 全局路径
# ======================================================================
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
LOG_ROOT = os.path.join(BASE_DIR, "logs")
MODEL_ROOT = os.path.join(BASE_DIR, "models")

# 中文字体：Windows 自带的微软雅黑 / 黑体
plt.rcParams["font.sans-serif"] = ["Microsoft YaHei", "SimHei", "DejaVu Sans"]
plt.rcParams["axes.unicode_minus"] = False

# ======================================================================
# 第 38–39 页：Taxi-v3 官方评估种子（PPT 原文：请勿修改 EVAL_SEED）
# 作用：让所有同学都在相同的出租车起始位置上评估，班级分数才可比。
# ======================================================================
EVAL_SEED_TAXI = [
    16, 54, 165, 177, 191, 191, 120, 80, 149, 178, 48, 38, 6, 125, 174, 73, 50,
    172, 100, 148, 146, 6, 25, 40, 68, 148, 49, 167, 9, 97, 164, 176, 61, 7, 54,
    55, 161, 131, 184, 51, 170, 12, 120, 113, 95, 126, 51, 98, 36, 135, 54, 82,
    45, 95, 89, 59, 95, 124, 9, 113, 58, 85, 51, 134, 121, 169, 105, 21, 30, 11,
    50, 65, 12, 43, 82, 145, 152, 97, 106, 55, 31, 85, 38, 112, 102, 168, 123, 97,
    21, 83, 158, 26, 80, 63, 5, 81, 32, 11, 28, 148,
]

# ======================================================================
# 环境 + 超参数配置（严格对应 PPT 第 34 页与第 39 页）
# ======================================================================
EXPERIMENTS = {
    # ---- 第一部分：冰冻湖 4x4，不滑（确定性） -------------------------
    "frozenlake4x4": dict(
        title="FrozenLake-v1（4x4，is_slippery=False）",
        env_id="FrozenLake-v1",
        env_kwargs=dict(map_name="4x4", is_slippery=False),
        n_training_episodes=10000,      # 第 34 页：总训练回合数
        learning_rate=0.7,              # 第 34 页：学习率
        gamma=0.95,                     # 第 34 页：折扣率
        max_steps=99,                   # 第 34 页：每回合最大步数
        max_epsilon=1.0,                # 第 34 页：开始时探索概率
        min_epsilon=0.05,               # 第 34 页：最小探索概率
        decay_rate=0.0005,              # 第 34 页：探索概率指数衰减率
        n_eval_episodes=100,            # 第 34 页：评估回合数
        eval_seed=None,                 # 第 34 页 eval_seed = []
        eval_every=1000,
        log_every=500,
    ),
    # ---- 第 31 页提到的另一种模式：4x4 + 滑冰（is_slippery=True） ------
    # 与第一部分只差一个变量 is_slippery，但冰面会以 1/3 概率把智能体推偏，
    # 所以同样的超参数明显更难学：这里把探索期拉长（min_epsilon 0.1 / decay 0.0003）
    # 并多训一倍回合，成功率才能从 ~20% 提到 ~75%。
    "frozenlake4x4_slippery": dict(
        title="FrozenLake-v1（4x4，is_slippery=True）",
        env_id="FrozenLake-v1",
        env_kwargs=dict(map_name="4x4", is_slippery=True),
        n_training_episodes=20000,
        learning_rate=0.7,
        gamma=0.95,
        max_steps=99,
        max_epsilon=1.0,
        min_epsilon=0.1,
        decay_rate=0.0003,
        n_eval_episodes=100,
        eval_seed=list(range(100)),
        eval_every=2000,
        log_every=1000,
    ),
    # ---- 第 31 页提到的进阶版：冰冻湖 8x8（地图变大） -----------------
    # 8x8 最优路径约 14 步，比 4x4 长一倍多；Q 表从 16x4 变成 64x4。
    # 注意：8x8 的稀疏奖励让"随机探索"很难碰到终点，如果沿用 4x4 的快速 epsilon 衰减，
    #       Q 表会一直是全 0（这一点在实验报告里作为对照实验记录）。
    #       所以这里把探索期拉长（min_epsilon 0.3 / decay 0.00003），让智能体先"逛遍"地图。
    "frozenlake8x8": dict(
        title="FrozenLake-v1（8x8，is_slippery=False）",
        env_id="FrozenLake-v1",
        env_kwargs=dict(map_name="8x8", is_slippery=False),
        n_training_episodes=60000,
        learning_rate=0.8,
        gamma=0.95,
        max_steps=200,
        max_epsilon=1.0,
        min_epsilon=0.3,
        decay_rate=0.00003,
        n_eval_episodes=100,
        eval_seed=list(range(100)),
        eval_every=6000,
        log_every=3000,
    ),
    # ---- 第二部分：Taxi-v3 -------------------------------------------
    "taxi": dict(
        title="Taxi-v3",
        env_id="Taxi-v3",
        env_kwargs=dict(),
        n_training_episodes=25000,      # 第 39 页：总训练回合数
        learning_rate=0.7,
        gamma=0.95,
        max_steps=99,
        max_epsilon=1.0,
        min_epsilon=0.05,
        decay_rate=0.005,               # 第 39 页：0.005（比冰冻湖衰减快，因为回合更长）
        n_eval_episodes=100,
        eval_seed=EVAL_SEED_TAXI,       # 第 39 页：官方评估种子
        eval_every=2500,
        log_every=1000,
    ),
}


# ======================================================================
# 工具：同时写控制台 + 日志文件（这就是"有日志输出的训练代码"）
# ======================================================================
class Tee:
    """把 print 的内容同时写进 stdout 和日志文件。"""

    def __init__(self, log_path):
        self.terminal = sys.stdout
        self.log_file = open(log_path, "w", encoding="utf-8")

    def write(self, message):
        self.terminal.write(message)
        self.log_file.write(message)
        self.log_file.flush()

    def flush(self):
        self.terminal.flush()
        self.log_file.flush()


def timestamp():
    return datetime.now().strftime("%H:%M:%S")


# ======================================================================
# 第 33 页：创建并初始化 Q 表
#   Q 表大小 = (状态数, 动作数)，初始全 0（"在开始时 Q 表是无用的"）
# ======================================================================
def initialize_q_table(state_space, action_space):
    return np.zeros((state_space, action_space))


# ======================================================================
# 第 33 页：贪婪策略（更新策略 / 最终策略）
#   pi*(s) = argmax_a Q(s,a)
# ======================================================================
def greedy_policy(q_table, state):
    return int(np.argmax(q_table[state]))


# ======================================================================
# 第 34 页：epsilon-greedy 策略（行为策略）
#   以 epsilon 的概率探索（随机动作），否则利用（贪心动作）
#   —— Q-Learning 是离线策略（off-policy）：行动用 epsilon-greedy，
#      更新时却用 max_a Q(s', a)，两套策略。
# ======================================================================
def epsilon_greedy_policy(q_table, state, epsilon, action_space):
    if random.random() > epsilon:
        return greedy_policy(q_table, state)          # 利用
    return int(action_space.sample())                 # 探索


# ======================================================================
# 第 35 页：训练循环
#   Q(s,a) <- Q(s,a) + lr * [ r + gamma * max_a' Q(s',a') - Q(s,a) ]
# ======================================================================
def train(
    env,
    q_table,
    n_training_episodes,
    max_steps,
    learning_rate,
    gamma,
    max_epsilon,
    min_epsilon,
    decay_rate,
    train_seed=0,
    log_every=500,
    eval_every=None,
    eval_kwargs=None,
):
    """
    返回：q_table、每回合数据、评估历史
    """
    episode_rewards = []
    episode_steps = []
    episode_success = []          # 是否到达目标（FrozenLake 的 reward==1 / Taxi 的 reward==20）
    epsilons = []
    eval_history = []             # [(episode, mean, std, success_rate), ...]

    started = time.time()
    eval_kwargs = eval_kwargs or {}

    for episode in range(n_training_episodes):
        # ---- epsilon 指数衰减：越往后越少探索 ----
        epsilon = min_epsilon + (max_epsilon - min_epsilon) * np.exp(-decay_rate * episode)
        epsilons.append(epsilon)

        # ---- 重置环境 ----
        # 每回合都用固定种子 reset(seed=train_seed+episode)：
        # 同一个 --seed 重跑，起点序列、滑冰随机性和探索随机性完全一致 -> 分数可复现。
        # 注意：gymnasium 的 env.reset(seed=...) 只给 env 播种，不给 action_space 播种，
        #       而 epsilon-greedy 的"探索"要用 env.action_space.sample()，
        #       所以这里再单独给 action_space 播种一次（这一行是可复现的关键）。
        if train_seed is not None:
            state, _ = env.reset(seed=train_seed + episode)
            env.action_space.seed(train_seed + episode)
        else:
            state, _ = env.reset()

        total_reward = 0.0
        reached_goal = False

        for step in range(max_steps):
            # 1) 用 epsilon-greedy 选择动作 At
            action = epsilon_greedy_policy(q_table, state, epsilon, env.action_space)

            # 2) 执行动作，观察 Rt+1 和 St+1
            new_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated

            # 3) Bellman 更新（Q-Learning 的核心这一行）
            #    注意：真实终止状态没有未来价值，next_q 取 0；
            #    PPT 截图里没做这个掩码，这里修正后结果更稳（FrozenLake 两者等价）。
            next_q = 0.0 if terminated else float(np.max(q_table[new_state]))
            q_table[state][action] += learning_rate * (
                reward + gamma * next_q - q_table[state][action]
            )

            total_reward += reward
            state = new_state
            if done:
                reached_goal = reward > 0
                break

        episode_rewards.append(total_reward)
        episode_steps.append(step + 1)
        episode_success.append(1 if reached_goal else 0)

        # ---- 每 log_every 回合打印一行进度（可直接粘进实验报告）----
        if (episode + 1) % log_every == 0 or episode == 0:
            window = episode_rewards[-min(100, len(episode_rewards)):]
            window_success = episode_success[-min(100, len(episode_success)):]
            window_steps = episode_steps[-min(100, len(episode_steps)):]
            elapsed_min = (time.time() - started) / 60.0
            print(
                f"[train] ep={episode + 1:>6}/{n_training_episodes} | "
                f"eps={epsilon:.3f} | "
                f"最近100回合平均回报={np.mean(window):6.2f} | "
                f"成功率={np.mean(window_success) * 100:5.1f}% | "
                f"平均步数={np.mean(window_steps):5.1f} | "
                f"Q非零项={np.count_nonzero(q_table):5d} | "
                f"用时={elapsed_min:5.2f}min"
            )

        # ---- 周期评估（用贪婪策略跑独立回合）----
        if eval_every and (episode + 1) % eval_every == 0:
            mean_r, std_r, succ = evaluate_agent(
                env, q_table, max_steps, eval_kwargs.get("n_eval", 20),
                eval_kwargs.get("eval_seed"),
            )
            eval_history.append((episode + 1, mean_r, std_r, succ))
            print(
                f"[eval ] ep={episode + 1:>6} | 贪婪策略 {eval_kwargs.get('n_eval', 20)} 局: "
                f"回报={mean_r:6.2f} ± {std_r:5.2f} | 成功率={succ * 100:5.1f}%"
            )

    return q_table, episode_rewards, episode_steps, episode_success, epsilons, eval_history


# ======================================================================
# 第 36 页：评估智能体（贪婪策略，返回平均回报 ± 标准差）
# ======================================================================
def evaluate_agent(env, q_table, max_steps, n_episodes, seeds=None):
    rewards = []
    successes = []
    lengths = []
    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_reward = 0.0
        for step in range(max_steps):
            action = greedy_policy(q_table, state)
            state, reward, terminated, truncated, _ = env.step(action)
            total_reward += reward
            if terminated or truncated:
                break
        rewards.append(total_reward)
        lengths.append(step + 1)
        successes.append(1 if total_reward > 0 else 0)
    return float(np.mean(rewards)), float(np.std(rewards)), float(np.mean(successes))


# ======================================================================
# 画训练曲线
# ======================================================================
def moving_average(values, window=100):
    values = np.asarray(values, dtype=float)
    if len(values) < window:
        return values
    kernel = np.ones(window) / window
    return np.convolve(values, kernel, mode="valid")


def plot_curve(run_dir, title, episode_rewards, epsilons, episode_success,
               eval_history, n_training_episodes):
    fig, axes = plt.subplots(2, 1, figsize=(12, 8), sharex=False)

    # 上图：每回合回报 + 100 回合滑动平均 + 周期评估
    ax = axes[0]
    ax.plot(np.arange(1, len(episode_rewards) + 1), episode_rewards,
            color="#93c5fd", linewidth=0.5, alpha=0.6, label="每回合回报")
    ma = moving_average(episode_rewards, 100)
    ax.plot(np.arange(100, 100 + len(ma)), ma,
            color="#1d4ed8", linewidth=2, label="100 回合滑动平均")
    if eval_history:
        xs = [h[0] for h in eval_history]
        ys = [h[1] for h in eval_history]
        errs = [h[2] for h in eval_history]
        ax.errorbar(xs, ys, yerr=errs, fmt="o--", color="#dc2626", markersize=4,
                    linewidth=1.2, capsize=3, label="周期评估（贪婪策略）")
    ax.set_xlabel("训练回合 episode")
    ax.set_ylabel("回报 reward")
    ax.set_title(f"{title} · 训练曲线（共 {n_training_episodes} 回合）")
    # 早期有几局极端负分（例如 Taxi 乱按上下客被 -10 罚到 -400），会把曲线压扁，
    # 这里按 0.5%~99.5% 分位数裁一下 y 轴，极端点仍在图外但趋势看得清。
    rewards_arr = np.asarray(episode_rewards, dtype=float)
    low = min(float(np.percentile(rewards_arr, 0.5)), 0.0) - 2
    high = max(float(np.percentile(rewards_arr, 99.5)), 1.0) + 1
    if high - low > 5:
        ax.set_ylim(low, high)
    ax.legend(loc="lower right", fontsize=9)
    ax.grid(alpha=0.3)

    # 下图：epsilon 衰减 + 探索/利用的阶段
    ax2 = axes[1]
    ax2.plot(np.arange(1, len(epsilons) + 1), epsilons, color="#f59e0b", linewidth=2,
             label="epsilon（探索概率）")
    ax2.fill_between(np.arange(1, len(epsilons) + 1), 0, epsilons, color="#f59e0b", alpha=0.15)
    ax2.set_xlabel("训练回合 episode")
    ax2.set_ylabel("epsilon")
    ax2.set_title("epsilon 指数衰减：从\"大胆探索\"到\"相信 Q 表\"")
    ax2.grid(alpha=0.3)
    ax2.legend(loc="upper right", fontsize=9)

    fig.tight_layout()
    path = os.path.join(run_dir, "train_curve.png")
    fig.savefig(path, dpi=130)
    plt.close(fig)
    return path


def write_progress_csv(run_dir, rewards, steps, success, epsilons, eval_history):
    path = os.path.join(run_dir, "progress.csv")
    with open(path, "w", newline="", encoding="utf-8-sig") as f:
        writer = csv.writer(f)
        writer.writerow(["episode", "epsilon", "episode_reward", "episode_steps", "success"])
        for i, (r, s, ok, eps) in enumerate(zip(rewards, steps, success, epsilons), start=1):
            writer.writerow([i, f"{eps:.6f}", f"{r:.6f}", s, ok])
        writer.writerow([])
        writer.writerow(["eval_episode", "eval_mean_reward", "eval_std_reward", "eval_success_rate"])
        for ep, mean_r, std_r, succ in eval_history:
            writer.writerow([ep, f"{mean_r:.6f}", f"{std_r:.6f}", f"{succ:.4f}"])
    return path


# ======================================================================
# 主流程
# ======================================================================
def main():
    parser = argparse.ArgumentParser(description="Q-Learning 训练（FrozenLake / Taxi）")
    parser.add_argument("--env", choices=list(EXPERIMENTS.keys()), default="frozenlake4x4",
                        help="用哪个环境训练")
    parser.add_argument("--episodes", type=int, default=None, help="覆盖总训练回合数")
    parser.add_argument("--learning-rate", type=float, default=None, help="覆盖学习率 lr")
    parser.add_argument("--gamma", type=float, default=None, help="覆盖折扣率 gamma")
    parser.add_argument("--decay-rate", type=float, default=None, help="覆盖 epsilon 衰减率")
    parser.add_argument("--min-epsilon", type=float, default=None, help="覆盖最小探索概率")
    parser.add_argument("--max-steps", type=int, default=None, help="覆盖每回合最大步数")
    parser.add_argument("--eval-episodes", type=int, default=None, help="覆盖最终评估回合数")
    parser.add_argument("--eval-every", type=int, default=None, help="每多少回合评估一次")
    parser.add_argument("--log-every", type=int, default=None, help="每多少回合打印一行日志")
    parser.add_argument("--seed", type=int, default=0, help="训练随机种子（默认 0）")
    parser.add_argument("--run-name", type=str, default=None, help="日志/模型文件名前缀")
    parser.add_argument("--slippery", action="store_true",
                        help="把 FrozenLake 的冰面切换成滑冰版（is_slippery=True，随机转移）")
    args = parser.parse_args()

    cfg = dict(EXPERIMENTS[args.env])
    if args.slippery:
        cfg["env_kwargs"] = dict(cfg["env_kwargs"], is_slippery=True)
        cfg["title"] += " · 滑冰版"
        # 滑冰版起点随机性大，固定 100 个起点种子才能公平比较
        cfg["eval_seed"] = list(range(100))
    # 命令行覆盖（方便做"一次只改一个变量"的对照实验）
    if args.episodes is not None:
        cfg["n_training_episodes"] = args.episodes
    if args.learning_rate is not None:
        cfg["learning_rate"] = args.learning_rate
    if args.gamma is not None:
        cfg["gamma"] = args.gamma
    if args.decay_rate is not None:
        cfg["decay_rate"] = args.decay_rate
    if args.min_epsilon is not None:
        cfg["min_epsilon"] = args.min_epsilon
    if args.max_steps is not None:
        cfg["max_steps"] = args.max_steps
    if args.eval_episodes is not None:
        cfg["n_eval_episodes"] = args.eval_episodes
    if args.eval_every is not None:
        cfg["eval_every"] = args.eval_every
    if args.log_every is not None:
        cfg["log_every"] = args.log_every
    cfg["train_seed"] = args.seed

    run_name = args.run_name or args.env
    run_dir = os.path.join(LOG_ROOT, run_name)
    os.makedirs(run_dir, exist_ok=True)
    os.makedirs(MODEL_ROOT, exist_ok=True)

    # 每次运行先设置随机种子：Python random 管探索、np 管 Q 表初始化
    random.seed(args.seed)
    np.random.seed(args.seed)

    tee = Tee(os.path.join(run_dir, "train_console.log"))
    sys.stdout = tee
    try:
        print("=" * 78)
        print(f"Q-Learning 训练 · {cfg['title']}")
        print(f"运行名称: {run_name}    开始时间: {datetime.now():%Y-%m-%d %H:%M:%S}")
        print("=" * 78)

        # ---------------- 第 31–32 页：创建并理解环境 ----------------
        env = gym.make(cfg["env_id"], render_mode="rgb_array", **cfg["env_kwargs"])
        state_space = env.observation_space.n
        action_space = env.action_space.n
        sample_state, _ = env.reset(seed=args.seed)

        print("\n[环境自检]")
        print(f"  env_id             = {cfg['env_id']}")
        print(f"  env_kwargs         = {cfg['env_kwargs']}")
        print(f"  观测空间           = {env.observation_space}  -> {state_space} 个状态")
        print(f"  动作空间           = {env.action_space}  -> {action_space} 个动作")
        if cfg["env_id"] == "FrozenLake-v1":
            ncol = env.unwrapped.ncol
            print(f"  一次 reset 的观测  = {sample_state}"
                  f"（整数 = 状态编号，{env.unwrapped.nrow}x{ncol} 地图按 row*{ncol}+col 编码）")
        else:
            print(f"  一次 reset 的观测  = {sample_state}（整数 = 状态编号，共 500 个）")
        print(f"  每回合最大步数     = {cfg['max_steps']}")

        # ---------------- 第 33 页：Q 表 ----------------
        q_table = initialize_q_table(state_space, action_space)
        print("\n[Q 表初始化]")
        print(f"  Q 表形状 = {q_table.shape}（= (状态数, 动作数)），初始值全为 0")
        print(f"  需要学习的参数个数 = {state_space * action_space}")

        # ---------------- 第 34 页：超参数 ----------------
        print("\n[超参数]")
        print(f"  n_training_episodes = {cfg['n_training_episodes']}")
        print(f"  learning_rate       = {cfg['learning_rate']}")
        print(f"  gamma               = {cfg['gamma']}")
        print(f"  max_steps           = {cfg['max_steps']}")
        print(f"  max_epsilon         = {cfg['max_epsilon']}")
        print(f"  min_epsilon         = {cfg['min_epsilon']}")
        print(f"  decay_rate          = {cfg['decay_rate']}")
        print(f"  n_eval_episodes     = {cfg['n_eval_episodes']}")
        print(f"  eval_seed           = {'官方 100 个种子' if cfg['eval_seed'] else '[]（不指定）'}")
        print(f"  train_seed          = {args.seed}")

        # ---------------- 第 35 页：训练 ----------------
        print(f"\n[开始训练] {timestamp()}  每 {cfg['log_every']} 回合打印一行，"
              f"每 {cfg['eval_every']} 回合评估一次")
        print(f"  可复现性：每回合 env.reset(seed={args.seed}+episode)，同 --seed 重跑分数一致")
        started = time.time()
        (q_table, rewards, steps, success, epsilons, eval_history) = train(
            env=env,
            q_table=q_table,
            n_training_episodes=cfg["n_training_episodes"],
            max_steps=cfg["max_steps"],
            learning_rate=cfg["learning_rate"],
            gamma=cfg["gamma"],
            max_epsilon=cfg["max_epsilon"],
            min_epsilon=cfg["min_epsilon"],
            decay_rate=cfg["decay_rate"],
            train_seed=args.seed,
            log_every=cfg["log_every"],
            eval_every=cfg["eval_every"],
            eval_kwargs=dict(
                n_eval=20,
                eval_seed=cfg["eval_seed"],
            ),
        )
        elapsed = (time.time() - started) / 60.0
        print(f"\n[训练完成] {timestamp()}  总用时 = {elapsed:.2f} 分钟")

        # ---------------- 第 36 页：最终评估 ----------------
        print(f"\n[最终评估] 用贪婪策略跑 {cfg['n_eval_episodes']} 局（独立于训练）")
        mean_r, std_r, succ = evaluate_agent(
            env, q_table, cfg["max_steps"], cfg["n_eval_episodes"], cfg["eval_seed"]
        )
        print(f"  Mean_reward = {mean_r:.2f} +/- {std_r:.2f}")
        print(f"  成功率      = {succ * 100:.0f}%")
        if eval_history:
            best = max(eval_history, key=lambda h: h[1])
            print(f"  训练途中最好的周期评估：{best[1]:.2f} ± {best[2]:.2f} @ 第 {best[0]} 回合")
        print(f"  Q 表非零项 = {np.count_nonzero(q_table)} / {q_table.size}")
        print(f"  Q 表最大价值 = {q_table.max():.4f}")

        # ---------------- 第 36–37 页：Q 表可视化 + 训练曲线 ----------------
        curve_path = plot_curve(run_dir, cfg["title"], rewards, epsilons, success,
                                eval_history, cfg["n_training_episodes"])
        csv_path = write_progress_csv(run_dir, rewards, steps, success, epsilons, eval_history)
        print(f"\n[输出] 训练曲线 = {os.path.relpath(curve_path, BASE_DIR)}")
        print(f"[输出] 进度 CSV = {os.path.relpath(csv_path, BASE_DIR)}")

        # ---------------- 第 37 页：保存模型（= "训练出的模型文件"） ----------------
        model = {
            "run_name": run_name,
            "title": cfg["title"],
            "env_id": cfg["env_id"],
            "env_kwargs": cfg["env_kwargs"],
            "max_steps": cfg["max_steps"],
            "n_training_episodes": cfg["n_training_episodes"],
            "n_eval_episodes": cfg["n_eval_episodes"],
            "eval_seed": cfg["eval_seed"],
            "learning_rate": cfg["learning_rate"],
            "gamma": cfg["gamma"],
            "max_epsilon": cfg["max_epsilon"],
            "min_epsilon": cfg["min_epsilon"],
            "decay_rate": cfg["decay_rate"],
            "train_seed": args.seed,
            "qtable": q_table,
            "metrics": {
                "mean_reward": mean_r,
                "std_reward": std_r,
                "success_rate": succ,
                "training_minutes": elapsed,
            },
            "trained_at": datetime.now().isoformat(timespec="seconds"),
        }
        model_path = os.path.join(MODEL_ROOT, f"qtable_{run_name}.pkl")
        with open(model_path, "wb") as f:
            pickle.dump(model, f)
        print(f"[输出] 模型文件 = {os.path.relpath(model_path, BASE_DIR)} "
              f"({os.path.getsize(model_path) / 1024:.1f} KB)")

        print("\n" + "=" * 78)
        print(f"完成。下一步：python evaluate_and_record.py --model models/qtable_{run_name}.pkl")
        print("=" * 78)
        env.close()
    finally:
        sys.stdout = tee.terminal
        tee.log_file.close()


if __name__ == "__main__":
    main()