"""
=============================================================================
CS486 · 第 9 讲模块：search_lab.py —— 把"LLM"放进搜索里：策略、价值函数、世界模型与几种搜索算法
=============================================================================

【这个模块要解决什么问题】
    第 8 讲把搜索交给经典规划器、把 LLM 放在搜索之外（LLM-Modulo）。这一讲反过来：让学出来的模型**参与**搜索。
    一个搜索过程需要四样东西（Tree of Thoughts、RAP、LATS、AlphaZero 都是它们的不同组合）：

        1. 状态转移（世界模型）：在状态 s 做动作 a 之后是什么状态？——可以是真实的模拟器，也可以是学出来的模型；
        2. 候选动作（策略 / 提议者）：在 s 有哪些值得试的动作？——第 8 讲的提议者给出每个动作的概率；
        3. 状态评估（价值函数 / 启发函数）：s 离目标还有多远？——可以是学出来的，也可以是人写的（misplaced）；
        4. 搜索算法：按什么顺序展开、保留哪些、什么时候回溯——贪心、束搜索、最佳优先、MCTS。

    本模块把这四样东西都做成可以替换的零件，三个脚本分别研究：
        代码 01：同一个提议者，放进不同的搜索算法里（真实模拟器 + 学出来的价值函数）；
        代码 02：价值函数的质量怎样决定搜索的成败——为什么"准确率"不够，误差落在哪些状态上更重要；
        代码 03：学出来的世界模型——一步的误差怎样累积，搜索怎样"钻"世界模型的空子，闭环怎样补救。

【复用】第 8 讲的 plan_lab（积木世界、验证器、A*、提议者），第 1 讲的 tiny_reasoner / hop_world，
    第 2 讲的 method_lab.fit_stream（训练循环），第 5 讲的 rl_lab.sample_fast（带 KV cache 的解码）。

【运行】
    python3 search_lab.py      # 自检（需要第 8 讲的提议者缓存；没有时先训练约 1 分钟）

【依赖】PyTorch。
"""

import functools
import heapq
import math
import os
import random
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
for _d in ("lecture08", "lecture05", "lecture04", "lecture02", "lecture01"):
    sys.path.insert(0, os.path.join(HERE, "..", _d))

import torch                      # noqa: E402

import hop_world as hw            # noqa: E402
import plan_lab as pl             # noqa: E402
import tiny_reasoner as tr        # noqa: E402

torch.set_num_threads(1)
CKPT_DIR = os.path.join(HERE, "checkpoints")


# ---------------------------------------------------------------------------
# 0. 一个小工具：一批字符串各自"最后一个位置"的 logits
# ---------------------------------------------------------------------------
@torch.no_grad()
def last_logits(model, texts):
    """
    把一批长短不一的字符串送进模型，取出每一条**最后一个字符处**的 logits（= 模型对"下一个字符"的打分）。

    【思路】右边补 PAD 拼成矩阵、做一次前向。注意力是因果的，右边的 PAD 不影响左边，所以每一行
        在自己最后一个真实字符处的输出，和单独送进去完全一样（第 1 讲 generate 的同一个道理）。
    【参数】model：TinyReasoner；texts：字符串列表（不能为空）
    【返回】张量 (len(texts), 词表大小)
    【PyTorch 小知识】lg[torch.arange(B), idx] 是"高级索引"：第 i 行取第 idx[i] 个位置。
    """
    seqs = [hw.encode(t) for t in texts]
    T = max(len(s) for s in seqs)
    X = torch.full((len(seqs), T), hw.PAD_ID, dtype=torch.long)
    for i, s in enumerate(seqs):
        X[i, :len(s)] = torch.tensor(s)
    lg = model(X)
    return lg[torch.arange(len(seqs)), torch.tensor([len(s) - 1 for s in seqs])]


def _load_or_train(name, cfg, batch_fn, quick, verbose, what):
    """
    读取（没有缓存就训练）一个 2 层 TinyReasoner。缓存文件 checkpoints/<name>.pt 里保存配置，读取时核对。

    【参数】cfg：{"layers", "steps", "seed", ...}；batch_fn：(step, rng) → [(输入, 期望输出), …]；what：打印用的说明
    【边界】quick 模式只训练 30 步、不读写缓存。
    """
    import method_lab as ml
    path = os.path.join(CKPT_DIR, f"{name}.pt")
    if not quick and os.path.exists(path):
        blob = torch.load(path, weights_only=False)
        if blob["config"] == cfg:
            m = tr.TinyReasoner(tr.ModelConfig(n_layers=cfg["layers"]))
            m.load_state_dict(blob["state_dict"])
            m.eval()
            return m
    tcfg = tr.TrainConfig(mode="cot", seed=cfg["seed"], steps=30 if quick else cfg["steps"])
    if verbose and not quick:
        print(f"  训练{what}（{cfg['layers']} 层，{cfg['steps']} 步）…", flush=True)
    model, secs = ml.fit_stream(tr.ModelConfig(n_layers=cfg["layers"]), tcfg, batch_fn)
    if not quick:
        os.makedirs(CKPT_DIR, exist_ok=True)
        torch.save({"config": cfg, "state_dict": model.state_dict(), "seconds": secs}, path)
        if verbose:
            print(f"  ⏱ 用时 {secs:.0f} 秒", flush=True)
    return model


# ---------------------------------------------------------------------------
# 1. 策略：把第 8 讲的提议者变成"每个合法动作的概率"
# ---------------------------------------------------------------------------
class Policy:
    """
    提议者当策略用：在状态 s（目标 g）下，每个**合法**动作 (x, y) 的概率。

    【思路】提议者的输入是 "<s>?<g>|"，它接着写的第一个动作就是"从 s 出发，它认为第一步该做什么"。
        P(x, y) = P(第一个字符 = x) · P(第二个字符 = y | 第一个字符 = x)。
        第二个因子需要把 x 接在输入后面再算一次，所以一个状态要算 1 + （最上面的积木数）条序列——放在一批里一次前向。
        只保留**合法**动作（真实模拟器告诉我们哪些合法），并把概率重新归一化；
        同时记下"模型放在合法动作上的总概率"（legal mass），它说明提议者自己有多"守规矩"。
    【用法】pol = Policy(model)；[(probs, mass), …] = pol(states, goal)；pol.calls 是累计的"模型调用"次数
        （每个不同的 (状态, 目标) 只算一次，结果缓存起来——搜索里同一个状态常常被问好几次）。
    【边界】状态等于目标时也会照常计算（搜索算法自己负责在目标处停下）。
    """

    def __init__(self, model):
        self.model = model
        self.cache = {}
        self.calls = 0

    def __call__(self, states, goal):
        todo = [s for s in dict.fromkeys(states) if (s, goal) not in self.cache]
        if todo:
            texts, spans = [], []
            for s in todo:
                base = pl.prompt(s, goal)
                tops = [t[-1] for t in s]
                spans.append((len(texts), tops))
                texts.append(base)
                texts += [base + x for x in tops]
            P = last_logits(self.model, texts).softmax(-1)
            for s, (i0, tops) in zip(todo, spans):
                self.calls += 1
                raw = {}
                for x, y in pl.legal_moves(s):
                    j = tops.index(x)
                    raw[(x, y)] = P[i0, hw.STOI[x]].item() * P[i0 + 1 + j, hw.STOI[y]].item()
                mass = sum(raw.values())
                self.cache[(s, goal)] = ({m: v / mass for m, v in raw.items()}, mass)
        return [self.cache[(s, goal)] for s in states]


def ranked_moves(policy, s, goal):
    """策略在 s 下的合法动作，按概率从大到小排列。"""
    probs, _ = policy([s], goal)[0]
    return sorted(probs, key=lambda m: -probs[m])


def _trace(prev, s):
    """从"父指针"字典 prev（状态 → (父状态, 动作) 或 None）回溯出到 s 的动作序列。"""
    plan = []
    while prev[s] is not None:
        ps, m = prev[s]
        plan.append(m)
        s = ps
    return plan[::-1]


# ---------------------------------------------------------------------------
# 2. 价值函数（"离目标还有几步"的估计）
# ---------------------------------------------------------------------------
VALUE = {"layers": 2, "steps": 3000, "n_min": 3, "n_max": 5, "seed": 0}
DIGIT_IDS = [hw.STOI[str(k)] for k in range(10)]


def train_value(quick=False, verbose=True):
    """
    读取（没有缓存就训练）学出来的价值函数：输入和提议者相同的 "<s>?<g>|"，输出最短计划的步数（一个数字，≥ 9 记作 9）。

    【思路】这对应 LLM 搜索里的"评估器"：ToT 让 LLM 给部分解打 sure / likely / impossible，
        AlphaZero / rStar-Math 训练一个价值网络或过程奖励模型。这里用和提议者**相同的训练分布**（3–5 块积木），
        标签是 A* 求出的真实距离——一个"有监督学出来的距离估计"，没有任何人写的规则。
    【返回】TinyReasoner（eval 模式）；缓存 checkpoints/value.pt
    """
    def batch_fn(step, rng):
        out = []
        for _ in range(64):
            init, goal, plan = pl.make_problem(rng, rng.randint(VALUE["n_min"], VALUE["n_max"]))
            out.append((pl.prompt(init, goal), str(min(len(plan), 9)) + "$"))
        return out
    return _load_or_train("value", VALUE, batch_fn, quick, verbose, "价值函数（3–5 块积木，标签是 A* 的最短步数）")


class ValueNet:
    """
    学出来的价值函数：返回"离目标还有几步"的期望值 Σ_k k·P(k)（连续值，越小越好）。

    【为什么用期望而不是 argmax】期望值能区分"有 60% 把握是 3 步"和"有 90% 把握是 3 步"，排序更细；
        argmax 只有 10 个取值，大量平局。
    【用法】V = ValueNet(model)；V(states, goal) → [float, …]；V.calls 是累计调用次数（同一状态只算一次）
    【边界】状态等于目标时直接返回 0（不问模型）。
    """

    def __init__(self, model):
        self.model = model
        self.cache = {}
        self.calls = 0

    def __call__(self, states, goal):
        todo = [s for s in dict.fromkeys(states) if (s, goal) not in self.cache and s != goal]
        if todo:
            P = last_logits(self.model, [pl.prompt(s, goal) for s in todo])[:, DIGIT_IDS].softmax(-1)
            ev = (P * torch.arange(10.0)).sum(-1).tolist()
            for s, v in zip(todo, ev):
                self.calls += 1
                self.cache[(s, goal)] = v
        return [0.0 if s == goal else self.cache[(s, goal)] for s in states]


@functools.lru_cache(maxsize=None)
def true_distance(s, goal):
    """真实的最短距离（A*，启发可采纳所以是最优）。【Python 小知识】lru_cache 把算过的结果记住，状态是元组、可以当字典键。"""
    return len(pl.astar(s, goal)[0])


def misplaced_value(states, goal):
    """人写的启发函数：不在最终位置上的积木数（第 8 讲 §2.3；可采纳，从不高估）。"""
    return [float(pl.misplaced(s, goal)) for s in states]


def _key_rng(seed, s, goal):
    """
    给 (种子, 状态, 目标) 一个**确定的**随机数生成器。

    【为什么不用缓存 + 一个共享的 rng】那样每个状态拿到的噪声取决于"它是第几个被问到的"——换一种搜索算法，
        同一个状态的噪声就变了。按键生成，同一个状态在任何算法里噪声都相同，比较才公平。
    【Python 小知识】random.Random(字符串) 用字符串的哈希（SHA-512）当种子，跨进程、跨运行都确定；
        内置的 hash() 对字符串每次运行都不同（PYTHONHASHSEED），不能用。
    """
    return random.Random(f"{seed}|{pl.state_str(s)}|{pl.state_str(goal)}")


def noisy_oracle(sigma, seed=0):
    """
    "带独立噪声的神谕"：真实距离 + N(0, σ²)，每个状态的噪声互相独立（但同一个状态每次都一样）。

    【用途】代码 02 演示 2：用 σ 控制价值函数的"准确率"，看搜索的成败怎样随准确率变化（Chen et al. 2024 的问题）。
    """
    def V(states, goal):
        return [0.0 if s == goal else true_distance(s, goal) + _key_rng(seed, s, goal).gauss(0, sigma)
                for s in states]
    return V


def shuffled_oracle(errors, seed=0):
    """
    "打乱了误差的神谕"：真实距离 + 从 errors（一组误差值）里随机抽一个。

    【用途】代码 02 演示 3 的对照：把学出来的价值函数的误差**原样**收集起来，再随机分给各个状态——
        误差的大小分布完全相同，只是"哪个状态错多少"和状态本身无关了。
        如果搜索在它上面成功、在学出来的价值函数上失败，问题就不在误差的大小，而在误差**落在哪些状态上**。
    """
    errors = list(errors)

    def V(states, goal):
        return [0.0 if s == goal else true_distance(s, goal) + _key_rng(seed, s, goal).choice(errors)
                for s in states]
    return V


# ---------------------------------------------------------------------------
# 3. 搜索算法（都用**真实的**模拟器 pl.apply_move；代码 03 换成学出来的世界模型）
# ---------------------------------------------------------------------------
def greedy_closed(policy, init, goal, max_steps, memory=True):
    """
    闭环贪心：每一步把**当前真实状态**交给提议者，执行它最想做的合法动作。

    【思路】和第 8 讲的"开环"不同：开环是提议者一口气写完整个计划（中途的状态全靠它自己在"脑子里"推演）；
        闭环是每步都由模拟器告诉它现在的状态，它只需要想"第一步"。这是最简单的"把模型放进搜索"——没有回溯。
    【参数】memory=True：记住走过的状态，不再走回去（选概率最高的、通向新状态的动作）；False：不记，可能来回打转
    【返回】(计划 或 None, 展开的状态数)
    【边界】最多 max_steps 步；memory=True 时所有动作都通向走过的状态就放弃。
    """
    s, plan, seen, exp = init, [], {init}, 0
    while s != goal and len(plan) < max_steps:
        exp += 1
        moves = ranked_moves(policy, s, goal)
        if memory:
            moves = [m for m in moves if pl.apply_move(s, m)[0] not in seen]
            if not moves:
                return None, exp
        s, _ = pl.apply_move(s, moves[0])
        seen.add(s)
        plan.append(moves[0])
    return (plan if s == goal else None), exp


def beam_search(policy, init, goal, width, max_depth, value=None, topb=None, stats=None):
    """
    束搜索：每一层只保留最好的 width 个状态。

    【两种打分】
        value=None：按策略的累计对数概率 Σ log P(动作)——"模型最相信的几条路"（只用策略，不用评估器）；
        value=函数：按价值函数估计的剩余距离（越小越好）——Tree of Thoughts 的 BFS：
                    每个状态提出候选（topb 个，None = 全部合法动作），评估器给所有候选打分，保留最好的 width 个。
    【参数】topb：每个状态只展开策略概率最高的 topb 个动作（ToT 里"提出 b 个想法"）；
           stats：传入一个字典时，记录每一层"保留下来的状态里有没有离目标最近的那个"（代码 02 用；要算真实距离，慢）
    【返回】(计划 或 None, 展开的状态数)
    【边界】一个状态一旦被**保留**过就不再保留第二次（避免兜圈子）；没被保留的候选以后还可以从别的路径进来。
        生成的子状态里出现目标就立即返回（不需要再评估）。
    """
    beam = [(0.0, init, [])]
    kept = {init}
    exp = 0
    for _ in range(max_depth):
        res = policy([b[1] for b in beam], goal)
        exp += len(beam)
        cand = []
        for (lp, s, plan), (probs, _) in zip(beam, res):
            moves = sorted(probs, key=lambda m: -probs[m])
            for m in (moves[:topb] if topb else moves):
                t, _ = pl.apply_move(s, m)
                if t == goal:
                    return plan + [m], exp
                if t not in kept:
                    cand.append((lp + math.log(probs[m] + 1e-12), t, plan + [m]))
        if not cand:
            break
        if value is None:
            order = sorted(range(len(cand)), key=lambda i: -cand[i][0])
        else:
            vs = value([c[1] for c in cand], goal)
            order = sorted(range(len(cand)), key=lambda i: vs[i])
        beam = []
        for i in order:
            if cand[i][1] in kept:
                continue
            kept.add(cand[i][1])
            beam.append(cand[i])
            if len(beam) == width:
                break
        if stats is not None:
            dist = [true_distance(c[1], goal) for c in cand]
            stats["layers"] = stats.get("layers", 0) + 1
            stats["kept_best"] = stats.get("kept_best", 0) + (min(true_distance(b[1], goal) for b in beam) == min(dist))
            if value is not None:
                stats.setdefault("errors", []).extend(v - d for v, d in zip(vs, dist))
    return None, exp


def best_first(policy, value, init, goal, budget):
    """
    最佳优先搜索：总是展开 f = 已走步数 + 价值估计 最小的状态（就是 A*，只是启发函数换成了学出来的价值函数）。

    【思路】和束搜索不同，它可以**回溯**到很早以前的分支：只要那里的 f 更小。ToT 的 DFS 版本也能回溯，但顺序不同。
        策略只用来决定同一个 f 时先放哪个子状态（平局时按策略概率的顺序）。
    【返回】(计划 或 None, 展开的状态数)；展开超过 budget 个状态就放弃
    【边界】价值函数不可采纳时，找到的计划不保证最短。
    """
    tie = 0
    heap = [(value([init], goal)[0], 0, tie, init)]
    prev, best_g, exp = {init: None}, {init: 0}, 0
    while heap and exp < budget:
        _, g, _, s = heapq.heappop(heap)
        if g > best_g[s]:
            continue
        if s == goal:
            return _trace(prev, s), exp
        exp += 1
        kids = []
        for m in ranked_moves(policy, s, goal):
            t, _ = pl.apply_move(s, m)
            if g + 1 < best_g.get(t, 1 << 30):
                best_g[t], prev[t] = g + 1, (s, m)
                kids.append(t)
        if kids:
            for t, v in zip(kids, value(kids, goal)):
                tie += 1
                heapq.heappush(heap, (best_g[t] + v, best_g[t], tie, t))
    return None, exp


class _Node:
    """MCTS 的树节点：状态、各动作的先验 P、访问次数 N、累计价值 W、子节点、节点自己的价值估计 v。"""
    __slots__ = ("s", "moves", "P", "N", "W", "kids", "v")

    def __init__(self, s):
        self.s = s
        self.kids = None


def mcts(policy, value, init, goal, budget, c=1.0, gamma=0.8):
    """
    蒙特卡洛树搜索（PUCT，AlphaZero 的选择公式）：

        选择：从根往下，每一步选 argmax_a  Q(s,a) + c · P(s,a) · √(N(s)+1) / (1 + N(s,a))
              （Q 是这个动作至今的平均价值；P 是策略的先验——策略越看好、试得越少，越值得试）；
        扩展：走到一个没展开过的节点，就问策略（得到 P）和价值函数（得到 v = γ^估计距离，离目标越近越接近 1）；
        回传：把 v 沿路径传回去，每往上一层乘一次 γ（越远的目标越不值钱）。
        还没试过的动作的 Q 用父节点的 v 代替（"first-play urgency"；用 0 会让搜索只往已经试过的方向钻）。

    【在规划里的特殊之处】模拟器是真实的、目标是可检查的，所以选择路径上一碰到目标就直接返回这条计划
        （不需要等访问次数收敛）。路径上重复出现的状态当作死路（价值 0）。
    【返回】(计划 或 None, 展开的节点数)
    【边界】路径深度最多 4·积木数；c、gamma 是可调参数（代码 01 的 --quick 之外只用默认值）。
    """
    max_depth = 4 * sum(len(t) for t in init)

    def expand(node):
        probs, _ = policy([node.s], goal)[0]
        node.moves = list(probs)
        node.P = [probs[m] for m in node.moves]
        node.N = [0] * len(node.moves)
        node.W = [0.0] * len(node.moves)
        node.kids = [None] * len(node.moves)
        node.v = gamma ** value([node.s], goal)[0]
        return node.v

    if init == goal:
        return [], 0
    root = _Node(init)
    expand(root)
    exp = 1
    while exp < budget:
        node, path, on_path = root, [], {init}
        while True:
            total = sum(node.N)
            best, bi = -1e18, 0
            for i in range(len(node.moves)):
                q = node.W[i] / node.N[i] if node.N[i] else node.v
                u = q + c * node.P[i] * math.sqrt(total + 1) / (1 + node.N[i])
                if u > best:
                    best, bi = u, i
            if node.kids[bi] is None:
                node.kids[bi] = _Node(pl.apply_move(node.s, node.moves[bi])[0])
            path.append((node, bi))
            node = node.kids[bi]
            if node.s == goal:
                return [p.moves[i] for p, i in path], exp
            if node.s in on_path or len(path) >= max_depth:
                v = 0.0
                break
            on_path.add(node.s)
            if node.kids is None:
                v = expand(node)
                exp += 1
                break
        for p, i in reversed(path):
            p.N[i] += 1
            p.W[i] += v
            v *= gamma
    return None, exp


# ---------------------------------------------------------------------------
# 4. 学出来的世界模型（代码 03）
# ---------------------------------------------------------------------------
WORLD = {"layers": 2, "steps": 3000, "n_min": 3, "n_max": 5, "seed": 0}


def wm_prompt(s, move):
    """世界模型的输入："<状态>?<动作>|"，例如 "ac;b?cb|"。期望输出：下一个状态 + "$"，例如 "a;bc$"。"""
    return pl.state_str(s) + "?" + move[0] + move[1] + "|"


def train_world_model(quick=False, verbose=True):
    """
    读取（没有缓存就训练）学出来的世界模型：给定状态和一个合法动作，预测下一个状态。

    【思路】这对应 RAP 里"让 LLM 预测动作的后果"：LLM 读了大量描述世界变化的文本，能写出"做了 a 之后世界变成什么样"，
        但不保证对。这里同样只在 3–5 块积木上训练（和提议者、价值函数相同）——它在 6、7 块上会怎样，是代码 03 的问题之一。
    【返回】TinyReasoner；缓存 checkpoints/world.pt
    """
    def batch_fn(step, rng):
        out = []
        for _ in range(64):
            n = rng.randint(WORLD["n_min"], WORLD["n_max"])
            s = pl.random_state(rng, "".join(sorted(rng.sample(pl.LETTERS, n))))
            m = rng.choice(pl.legal_moves(s))
            out.append((wm_prompt(s, m), pl.state_str(pl.apply_move(s, m)[0]) + "$"))
        return out
    return _load_or_train("world", WORLD, batch_fn, quick, verbose, "世界模型（3–5 块积木：状态 + 动作 → 下一个状态）")


class WorldModel:
    """
    学出来的世界模型：WM([(s, move), …]) → [预测的下一个状态 或 None, …]。

    【思路】贪心解码（第 5 讲带 KV cache 的 sample_fast，温度 0），截到 "$"，解析成状态。
        输出里出现非法字符（不是字母或 ";"）时返回 None——"模型说不出一个像样的状态"。
    【用法】WM.calls 是累计调用次数（同一个 (状态, 动作) 只算一次）。
    【边界】预测的状态可能"丢了"积木或者多出积木——它只是一串字符，不受物理规律约束；搜索要能处理这种状态。
    """

    def __init__(self, model):
        self.model = model
        self.cache = {}
        self.calls = 0

    def __call__(self, pairs):
        import rl_lab as rl
        todo = [p for p in dict.fromkeys(pairs) if p not in self.cache]
        if todo:
            outs = rl.sample_fast(self.model, [wm_prompt(s, m) for s, m in todo], 2 * len(pl.LETTERS), 0, None)
            for p, o in zip(todo, outs):
                self.calls += 1
                txt = o.split("$", 1)[0]
                ok = bool(txt) and all(c in pl.LETTERS + ";" for c in txt)
                self.cache[p] = pl.parse_state(txt) if ok else None
        return [self.cache[p] for p in pairs]


def astar_in_model(wm, init, goal, max_expand):
    """
    在学出来的世界模型**里面**做 A*（启发函数 misplaced）：后继状态全部由 wm 预测，不碰真实模拟器。

    【思路】这就是"用世界模型规划"（RAP、MuZero 的思路）：只要模型说"做完这些动作就到目标了"，搜索就交出这个计划。
        计划在真实世界里能不能执行，要另外用可靠的验证器检查——代码 03 的核心问题。
    【返回】(计划 或 None, 展开的状态数)
    【边界】合法动作是根据**模型以为的**当前状态算的；模型预测出 None 的分支直接丢掉。
    """
    tie = 0
    heap = [(pl.misplaced(init, goal), 0, tie, init)]
    prev, best_g, exp = {init: None}, {init: 0}, 0
    while heap and exp < max_expand:
        _, g, _, s = heapq.heappop(heap)
        if g > best_g[s]:
            continue
        if s == goal:
            return _trace(prev, s), exp
        exp += 1
        moves = pl.legal_moves(s)
        for m, t in zip(moves, wm([(s, m) for m in moves])):
            if t is None or g + 1 >= best_g.get(t, 1 << 30):
                continue
            best_g[t], prev[t] = g + 1, (s, m)
            tie += 1
            heapq.heappush(heap, (g + 1 + pl.misplaced(t, goal), g + 1, tie, t))
    return None, exp


def rollout_in_model(wm, s, moves):
    """
    在世界模型里依次执行一串动作，返回模型以为的最终状态。

    【返回】状态；某一步在模型以为的状态里不合法、或者模型说不出状态时返回 None
    """
    for m in moves:
        if s is None or m not in pl.legal_moves(s):
            return None
        s = wm([(s, m)])[0]
    return s


def mpc(wm, init, goal, max_steps, max_expand):
    """
    闭环（模型预测控制，MPC）：在世界模型里规划 → 只在真实世界里执行**第一步** → 观察真实的新状态 → 重新规划。

    【思路】世界模型的错误在开环执行时会一直累积（第 k 步之后的模拟全建立在错误的状态上）；
        闭环每一步都用真实观察把"模型以为的状态"纠正回来，错误只影响眼前这一步的决策。
        这是智能体在环境里行动的基本形式（第 11 讲），也是第 10 讲"运行时重规划"的最简单版本。
    【返回】(是否到达目标, 实际走了几步, 规划了几次)
    【边界】某次规划失败、或者第一步在真实世界里不合法，就算失败。
    """
    s, steps, replans = init, 0, 0
    while s != goal and steps < max_steps:
        plan, _ = astar_in_model(wm, s, goal, max_expand)
        replans += 1
        if not plan:
            return False, steps, replans
        s, err = pl.apply_move(s, plan[0])
        if err:
            return False, steps, replans
        steps += 1
    return s == goal, steps, replans


# ---------------------------------------------------------------------------
# 自检
# ---------------------------------------------------------------------------
def _self_test():
    model = pl.train_proposer()
    pol = Policy(model)
    init, goal = pl.parse_state("ac;b"), pl.parse_state("abc")
    (probs, mass), = pol([init], goal)
    assert set(probs) == set(pl.legal_moves(init)) and abs(sum(probs.values()) - 1) < 1e-6 and 0 < mass <= 1
    best = max(probs, key=probs.get)
    print(f"策略：状态 'ac;b' → 目标 'abc'，合法动作 {len(probs)} 个，最可能 {best[0] + best[1]}（{probs[best]:.2f}），"
          f"模型放在合法动作上的总概率 {mass:.2f}")

    # 用真实距离当价值函数时，几种搜索都应该能解出小题，而且束搜索 / 最佳优先找到的计划合法
    oracle = noisy_oracle(0.0)
    rng = random.Random(1)
    for _ in range(10):
        a, b, opt = pl.make_problem(rng, 4)
        for plan, _ in (beam_search(pol, a, b, 4, 12, value=oracle), best_first(pol, oracle, a, b, 200),
                        mcts(pol, oracle, a, b, 400)):
            assert plan is not None and pl.validate(a, b, pl.plan_to_str(plan)).ok
        plan, _ = best_first(pol, oracle, a, b, 200)
        assert len(plan) == len(opt)                     # 价值 = 真实距离（可采纳）时，最佳优先就是 A*，计划最短
    print("真实距离当价值函数：束搜索、最佳优先、MCTS 在 10 道 4 块积木的题上都找到合法计划；最佳优先的计划最短")

    # 确定的噪声：同一个状态在不同调用里噪声相同
    v1, v2 = noisy_oracle(1.0, seed=3), noisy_oracle(1.0, seed=3)
    assert v1([init], goal) == v2([init], goal) and v1([goal], goal) == [0.0]
    sh = shuffled_oracle([5.0], seed=0)
    assert sh([init], goal)[0] == true_distance(init, goal) + 5.0
    print(f"带噪声的神谕：同一个状态两次调用结果相同（{v1([init], goal)[0]:.3f}）；真实距离 {true_distance(init, goal)}")

    # 世界模型的接口（quick：只训练 30 步，输出多半不对，只检查不崩溃）
    wmq = WorldModel(train_world_model(quick=True, verbose=False))
    out = wmq([(init, ("c", "b")), (init, ("c", "_"))])
    assert len(out) == 2 and all(o is None or isinstance(o, tuple) for o in out)
    assert rollout_in_model(lambda pairs: [pl.apply_move(s, m)[0] for s, m in pairs], init, [("c", "_"), ("b", "a"), ("c", "b")]) == goal
    ok, steps, _ = mpc(lambda pairs: [pl.apply_move(s, m)[0] for s, m in pairs], init, goal, 10, 100)
    assert ok and steps == 3
    print("世界模型：接口正常；用真实模拟器代替世界模型时，rollout 和 MPC 都在 3 步内到达目标")
    print("自检通过。")


if __name__ == "__main__":
    _self_test()