"""
=============================================================================
CS486 · 第 8 讲 · 代码 03：LLM-Modulo——让模型提议、让验证器把关
=============================================================================

【这个脚本要回答的问题】（讲义 §5）
    Kambhampati et al.（2024，"LLMs Can't Plan, But Can Help Planning in LLM-Modulo Frameworks"）主张：
    不要指望 LLM 自己做出正确的计划，也不要指望它检查自己（第 7 讲），而是把它放进一个"生成–测试–批评"的循环：
    LLM 提出候选计划，**可靠的**外部验证器检查它、指出错在哪里，LLM 根据反馈修改，直到通过验证。
    Valmeekam, Marquez & Kambhampati（2023，"Can Large Language Models Really Improve by Self-critiquing Their Own Plans?"）
    在 100 道积木世界题上：GPT-4 单独做对 40 道；让 GPT-4 自己当批评者、回溯提示（backprompting），55 道；
    换成可靠的验证器 VAL 当批评者，88 道。

    我们用代码 02 的提议者做同样的事（plan_lab.modulo_solve）：
        验证器说"第 i 步的 x 不在最上面" → 保留前 i 步，第 i 步换掉 x，接着写；
        验证器说"每步都合法但没到目标"   → 从当前状态出发，把剩下的计划接着写完。

    演示 1：反馈的轮数 vs 只采样不反馈（生成–测试）——各能解决多少题？用了多少次提议者？
    演示 2：批评者不可靠时会怎样？——可靠的验证器 vs 会漏判的批评者 vs 提议者的自我批评。
    演示 3：用第 7 讲的四参数模型看第一轮：为什么可靠的验证器只会帮忙、不会帮倒忙？
    演示 4：成本——提议者的调用次数 vs A* 展开的状态数。

【运行】
    python3 03_llm_modulo.py           # 约 1 分钟（需要代码 02 的提议者缓存；没有时先训练约 1 分钟）
    python3 03_llm_modulo.py --quick   # 冒烟测试

【依赖】PyTorch + 本目录的 plan_lab.py。
"""

import argparse
import os
import sys
import time

import torch

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

import plan_lab as pl                                                    # noqa: E402
from console import pct, print_header, print_observations, print_table  # noqa: E402

NS = (4, 5, 6, 7)
ROUNDS = (0, 1, 3, 10)
SAMPLES = (1, 4, 10, 30)


def sample_and_test(model, probs, N, critic, seed=0):
    """
    生成–测试：每道题最多采样 N 个计划（温度 1.0），第一个被批评者接受的就作为答案；都不被接受时用最后一个。

    【返回】(真正合法的比例, 被批评者接受的比例, 平均采样次数)
    【边界】"真正合法"始终用可靠的验证器判定，和批评者是谁无关。
    """
    gen = torch.Generator().manual_seed(seed)
    n_blocks = len({b for t in probs[0][0] for b in t})
    prompts = [pl.prompt(a, b) for a, b, _ in probs]
    chosen = [None] * len(probs)
    accepted = [False] * len(probs)
    used = [0] * len(probs)
    for _ in range(N):
        todo = [i for i in range(len(probs)) if not accepted[i]]
        if not todo:
            break
        plans = pl.propose_sampled(model, [prompts[i] for i in todo], 2 * n_blocks + 2, 1.0, gen)
        for i, s in zip(todo, plans):
            used[i] += 1
            chosen[i] = s
            if critic(probs[i][0], probs[i][1], s).ok:
                accepted[i] = True
    true_ok = sum(pl.validate(a, b, s).ok for (a, b, _), s in zip(probs, chosen)) / len(probs)
    return true_ok, sum(accepted) / len(probs), sum(used) / len(probs)


def self_critic(model, threshold):
    """
    提议者的自我批评：它给这个计划的对数概率 ≥ threshold 就"接受"（第 7 讲：内在反馈）。
    自我批评说不出错在哪一步，所以只用在"生成–测试"里（不接受就重新采样）。
    """
    def critic(init, goal, plan_s):
        lp = pl.plan_logprob(model, [pl.prompt(init, goal)], [plan_s])[0]
        return pl.Verdict(lp >= threshold, "ok" if lp >= threshold else "self", 0, "", init)
    return critic


def main():
    ap = argparse.ArgumentParser(description="LLM-Modulo")
    ap.add_argument("--quick", action="store_true", help="冒烟测试")
    args = ap.parse_args()
    t0 = time.perf_counter()
    count = 10 if args.quick else 100

    print_header("准备：代码 02 的提议者；每个 n 用 %d 道题（和代码 02 的测试集不同的种子）" % count)
    model = pl.train_proposer(quick=args.quick)
    sets = {n: pl.test_problems(n, count, seed=1) for n in NS}

    print_header("演示 1：验证器反馈的轮数 vs 只采样、不反馈（生成–测试）")
    rows, mod, gt = [], {}, {}
    for n in NS:
        cells = [n]
        for R in ROUNDS:
            res = [pl.modulo_solve(model, a, b, R, 2 * n + 2) for a, b, _ in sets[n]]
            mod[(n, R)] = (sum(r[2] for r in res) / count, sum(r[1] for r in res) / count)
            cells.append(pct(mod[(n, R)][0], 0))
        for N in SAMPLES:
            ok, _, used = sample_and_test(model, sets[n], N, pl.validate)
            gt[(n, N)] = (ok, used)
            cells.append(pct(ok, 0))
        rows.append(cells)
    print_table(["n", *[f"反馈 {R} 轮" for R in ROUNDS], *[f"采样 {N} 次" for N in SAMPLES]],
                rows, [4] + [10] * len(ROUNDS) + [10] * len(SAMPLES))
    print("  （反馈 0 轮 = 提议者直接给出的计划；采样 N 次 = 最多采样 N 个计划，取第一个通过验证的）")
    rows = [[n, f"{1 + mod[(n, 10)][1]:.1f}", f"{gt[(n, 30)][1]:.1f}"] for n in NS]
    print_table(["n", "反馈 10 轮：平均提议次数", "采样 30 次：平均采样次数"], rows, [4, 26, 26])

    print_header("演示 2：批评者不可靠时（n = 5；'声称成功' = 批评者接受了，'真正成功' = 可靠的验证器判定合法）")
    probs5 = sets[5]
    dev = pl.test_problems(5, count, seed=2)
    dev_plans = pl.propose(model, [pl.prompt(a, b) for a, b, _ in dev], 12)
    dev_lp = pl.plan_logprob(model, [pl.prompt(a, b) for a, b, _ in dev], dev_plans)
    dev_ok = [pl.validate(a, b, s).ok for (a, b, _), s in zip(dev, dev_plans)]
    best, tau = -1.0, 0.0
    for t in sorted(set(dev_lp)):
        tpr = sum(1 for x, y in zip(dev_lp, dev_ok) if y and x >= t) / max(1, sum(dev_ok))
        fpr = sum(1 for x, y in zip(dev_lp, dev_ok) if not y and x >= t) / max(1, len(dev_ok) - sum(dev_ok))
        if tpr - fpr > best:
            best, tau = tpr - fpr, t
    critics = {
        "可靠的验证器": pl.validate,
        "漏判 30% 的批评者": pl.noisy_critic(0.3, seed=1),
        "漏判 60% 的批评者": pl.noisy_critic(0.6, seed=1),
    }
    rows, crit = [], {}
    for name, c in critics.items():
        res = [pl.modulo_solve(model, a, b, 10, 12, critic=c) for a, b, _ in probs5]
        claimed = sum(r[3] for r in res) / count
        true = sum(r[2] for r in res) / count
        crit[name] = (claimed, true)
        rows.append([name, "反馈 10 轮", pct(claimed), pct(true)])
    for name, c in (("可靠的验证器", pl.validate), ("提议者的自我批评", self_critic(model, tau))):
        true, claimed, _ = sample_and_test(model, probs5, 10, c, seed=3)
        crit[name + "（采样）"] = (claimed, true)
        rows.append([name, "采样 10 次", pct(claimed), pct(true)])
    print_table(["批评者", "做法", "声称成功", "真正成功"], rows, [20, 12, 10, 10])
    print(f"  （自我批评：计划的对数概率 ≥ {tau:.3f} 就接受，阈值在另一组 n = 5 的题上按 Youden 指数选）")

    print_header("演示 3：第一轮反馈的四个参数（第 7 讲的模型），n = 4–7 合在一起")
    flagged = fixed = broke = n_wrong = n_right = 0
    for n in NS:
        for a, b, _ in sets[n]:
            p0 = pl.propose(model, [pl.prompt(a, b)], 2 * n + 2)[0]
            ok0 = pl.validate(a, b, p0).ok
            _, _, ok1, _ = pl.modulo_solve(model, a, b, 1, 2 * n + 2)
            if ok0:
                n_right += 1
                broke += not ok1
            else:
                n_wrong += 1
                flagged += 1
                fixed += ok1
    tpr = flagged / max(1, n_wrong)
    print_table(["TPR", "FPR", "修复率", "改坏率", "错的计划 / 对的计划"],
                [[f"{tpr:.2f}", "0.00", f"{fixed / max(1, flagged):.2f}", f"{broke / max(1, n_right):.2f}",
                  f"{n_wrong} / {n_right}"]], [8, 8, 10, 10, 20])
    p_fix1 = fixed / max(1, flagged)

    print_header("演示 4：成本——提议者调用次数 vs A* 展开的状态数（每个 n 的平均）")
    rows = []
    for n in NS:
        a_nodes = sum(pl.astar(a, b)[1] for a, b, _ in sets[n]) / count
        rows.append([n, pct(mod[(n, 10)][0], 0), f"{1 + mod[(n, 10)][1]:.1f}", "100%", f"{a_nodes:.1f}"])
    print_table(["n", "LLM-Modulo（10 轮）成功", "平均提议次数", "A* 成功", "A* 平均展开"], rows, [4, 24, 14, 10, 14])

    if args.quick:
        print("\n（--quick 冒烟测试：数字没有意义，不打印观察。）")
        print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")
        return
    sv, n30, n60 = crit["可靠的验证器"], crit["漏判 30% 的批评者"], crit["漏判 60% 的批评者"]
    selfc = crit["提议者的自我批评（采样）"]
    print_observations([
        f"可靠的验证器 + 反馈让同一个提议者明显变好（演示 1）：n = 5 时从 {pct(mod[(5, 0)][0], 0)}（直接给出）到 "
        f"{pct(mod[(5, 10)][0], 0)}（10 轮反馈），n = 4 时从 {pct(mod[(4, 0)][0], 0)} 到 {pct(mod[(4, 10)][0], 0)}。"
        "这就是 LLM-Modulo 的核心主张：LLM 自己做不出正确的计划，但它能提出'离正确不远'的候选，验证器负责把关。"
        "而且通过验证的计划**保证**是对的——这是单靠 LLM 永远给不了的。",
        f"反馈比盲目重采样更省调用（演示 1）：n = 5 时 10 轮反馈平均调用提议者 {1 + mod[(5, 10)][1]:.1f} 次、成功 "
        f"{pct(mod[(5, 10)][0], 0)}；只采样不反馈，最多 10 次（平均 {gt[(5, 10)][1]:.1f} 次）成功 {pct(gt[(5, 10)][0], 0)}，"
        f"最多 30 次（平均 {gt[(5, 30)][1]:.1f} 次）才到 {pct(gt[(5, 30)][0], 0)}。定位反馈（错在第几步、为什么）把搜索引向正确的方向"
        "（第 7 讲：定位反馈有效、二值反馈要靠运气）；但采样的多样性在预算大时有更高的上限。",
        f"提议者离得太远时，验证器也救不回来（演示 1）：n = 7 时 10 轮反馈只到 {pct(mod[(7, 10)][0], 0)}、30 次采样 "
        f"{pct(gt[(7, 30)][0], 0)}。LLM-Modulo 的效果取决于提议者能不能提出'在几步修改之内可达'的候选——"
        f"第一轮反馈的修复率只有 {p_fix1:.2f}（演示 3）。验证器保证了正确性，不保证能找到解。",
        f"批评者必须可靠（演示 2，n = 5）：漏判 30% 的批评者让系统'声称'成功 {pct(n30[0])}，真正成功只有 {pct(n30[1])}；"
        f"漏判 60% 时声称 {pct(n60[0])}、真正 {pct(n60[1])}；可靠的验证器声称和真正都是 {pct(sv[1])}。"
        f"用提议者自己的把握当批评者（第 7 讲的内在反馈）更糟：采样 10 次、只接受它自己有把握的计划，真正成功 {pct(selfc[1])}——"
        f"还不如直接贪心（{pct(mod[(5, 0)][0], 0)}）。**LLM-Modulo 的'Modulo'指的是外部的、可靠的验证器**，换成模型自己就失去了全部保证。",
        "第 7 讲的四参数模型解释了为什么可靠的验证器只会帮忙（演示 3）：它的 FPR = 0（从不把合法的计划判为不合法），"
        "所以改坏率无从发生，一轮反馈的收益 = (1 − a)·TPR·修复率 ≥ 0。第 7 讲里 PRM-full 的误报把对的改错，在这里不会出现。",
        f"但别忘了基线（演示 4）：A* 在所有题上 100% 成功、找到最短计划，平均只展开 {sum(pl.astar(a, b)[1] for a, b, _ in sets[7]) / count:.1f} 个状态"
        "（n = 7）——比调用一次提议者便宜得多。在规则完全已知、能写成 PDDL 的问题上，正确的做法是用规划器；"
        "LLM 的价值在于规则不完全已知、需要常识或自然语言理解的地方——例如把自然语言描述翻译成 PDDL（LLM+P，第 9 讲），"
        "或者在验证器只能部分检查的开放问题上（第 10 讲的旅行规划）。",
    ])
    print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")


if __name__ == "__main__":
    main()