"""
=============================================================================
第 5 讲 · 02：ADaPT —— 做不了的时候才分解（as-needed decomposition）
=============================================================================

对应讲义：lectures/lecture05_reasoning_planning.md 第 3 节（Prasad et al. 2024）

-----------------------------------------------------------------------------
【设定】
-----------------------------------------------------------------------------
任务：第 1 讲的 k 步算术链（k = 2 ~ 10）。
executor（执行者）：一次调用就要交出结果，**不写中间步骤**（第 1 讲的 direct 模式）。
    一次调用能可靠完成的串行步数有限：1B 模型约 2 步，100B 模型约 4 步（toy_llm.depth_capacity）。
    超过这个上限的部分几乎只能靠猜 —— 这对应"任务太复杂，executor 一口气做不完"。
planner（规划者）：把一个任务拆成两个子任务（前一半、后一半，必须都完成：AND）。每次拆分有 2% 的概率出错
    （例如漏掉一步），这代表"分解本身也会引入错误"。
环境的成功信号：每做完一个（子）任务，环境告诉我们成功没有，有一点噪声（TPR 0.98，FPR 0.02）——
    对应 ALFWorld 里"环境告诉你东西放进去了没有"。环境检查的是"过程是否真的做对"，不是"数字碰巧对了"。

三种策略：
    iterative executor：不规划，executor 直接做整个任务，失败了重试，最多 3 次。
    plan-and-execute  ：先把任务固定地拆成每块 2 步的子任务，每块最多尝试 2 次，任何一块失败整个任务就失败。
    ADaPT             ：先让 executor 直接做（最多尝试 2 次）；还是失败，才让 planner 把它拆成两半，
                        对每一半**递归地**重复这个过程，最多拆 3 层。
                        任务简单就不拆，任务难就多拆；executor 强就少拆，弱就多拆。

运行时间：约 1 秒（只用标准库）。
"""

import os
import random
import sys

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

from toy_llm import Problem, ToyLLM, make_problem, mean, pad  # noqa: E402

TEMPERATURE = 0.7
ENV_TPR, ENV_FPR = 0.98, 0.02
ATTEMPTS = 2          # 每个（子）任务最多尝试几次（三种策略相同）
PLAN_ERROR = 0.02
MAX_DEPTH = 3


class Counter:
    """记录 LLM 调用次数（executor 和 planner 分开计）。"""

    def __init__(self):
        self.executor = 0
        self.planner = 0


def run_executor(executor, start, steps, rng, counter):
    """
    executor 一次调用：从 start 出发，一口气做完 steps 这些运算（direct 模式）。

    【思路】构造一个子问题交给 ToyLLM(cot=False)。步骤下标从 0 重新开始数：每次调用都是一次新的前向传播，
        串行上限对每次调用分别生效 —— 这正是"拆小了就能做"的原因。
    【返回】(结果, 过程是否真的做对)。"过程做对" = 每一步都算对（Sample.local_ok 全为 True）。
    【边界】为什么不只看结果？toy world 里加法的 bug（少 10）和减法的 bug（多 10）会互相抵消，
        direct 模式下 bug 很多，"答案碰巧对了"非常常见（第 3 讲演示 2）。ALFWorld 这类任务里，
        环境检查的是真实的状态，不存在"碰巧成功"，所以这里也按过程判断。
    """
    counter.executor += 1
    s = executor.generate(Problem(-1, start, tuple(steps)), rng, TEMPERATURE)
    return s.answer, all(s.local_ok)


def env_says_success(ok, rng):
    """环境的成功信号：真的做对时以 TPR 报告成功，没做对时以 FPR 误报成功。"""
    return rng.random() < (ENV_TPR if ok else ENV_FPR)


def plan_split(steps, rng, counter):
    """
    planner 把任务拆成两半。2% 的概率出错：漏掉一步（错误会传到最终结果里）。

    【返回】(前一半, 后一半)
    """
    counter.planner += 1
    mid = len(steps) // 2
    first, second = list(steps[:mid]), list(steps[mid:])
    if rng.random() < PLAN_ERROR:
        if len(second) > 1:
            second = second[1:]
        elif len(first) > 1:
            first = first[1:]
    return first, second


def iterative_executor(executor, problem, rng, counter, attempts=3):
    """不规划：直接做整个任务，环境说失败就重试，最多 attempts 次。返回 (最终交出的结果, 过程是否做对)。"""
    result, ok = None, False
    for _ in range(attempts):
        result, ok = run_executor(executor, problem.start, problem.steps, rng, counter)
        if env_says_success(ok, rng):
            break
    return result, ok


def plan_and_execute(executor, problem, rng, counter, chunk=2, attempts=ATTEMPTS):
    """
    先固定地拆成每块 chunk 步的子任务（planner 只调用一次），再依次执行，每块最多尝试 attempts 次。

    【边界】某一块所有尝试都失败时，整个任务失败 —— 计划是固定的，不会因为执行失败而调整。
    """
    counter.planner += 1
    x, all_ok = problem.start, True
    for j in range(0, problem.n_steps, chunk):
        block = list(problem.steps[j:j + chunk])
        accepted = False
        for _ in range(attempts):
            y, ok = run_executor(executor, x, block, rng, counter)
            if env_says_success(ok, rng):
                accepted = True
                break
        if not accepted:
            return None, False
        x, all_ok = y, all_ok and ok
    return x, all_ok


def adapt(executor, start, steps, rng, counter, depth=0, depths=None):
    """
    ADaPT：先直接做；失败了才拆成两半，递归地处理每一半。

    【思路】
        1. 让 executor 直接做这个（子）任务（最多 ATTEMPTS 次）；环境说成功就返回结果。
        2. 否则，如果还没到最大深度、而且任务多于 1 步：让 planner 拆成两半（AND），
           先递归地做前一半，拿它的结果作为后一半的起点，再递归地做后一半。
        3. 到了最大深度还失败：交出 executor 的结果（可能是错的）。
        这样，**分解的粒度由 executor 的能力和任务的难度共同决定**：executor 能一口气做完的，就不拆。
    【参数】start、steps：（子）任务；depth：当前的递归深度；depths：记录用到的深度（list，可选）
    【返回】(结果, 过程是否做对)。注意 planner 拆错（漏掉一步）时，每个子任务都可能"做对"，最终结果却是错的。
    """
    if depths is not None:
        depths.append(depth)
    for _ in range(ATTEMPTS):
        result, ok = run_executor(executor, start, steps, rng, counter)
        if env_says_success(ok, rng):
            return result, ok
    if depth >= MAX_DEPTH or len(steps) <= 1:
        return result, ok
    first, second = plan_split(steps, rng, counter)
    mid, ok1 = adapt(executor, start, first, rng, counter, depth + 1, depths)
    end, ok2 = adapt(executor, mid, second, rng, counter, depth + 1, depths)
    return end, ok1 and ok2


def main():
    rng = random.Random(130)
    lengths = (2, 4, 6, 8, 10)
    problems = {k: [make_problem(random.Random(1300 + k), k, pid=i) for i in range(150)] for k in lengths}
    # 【Python 小知识】make_problem 用同一个 rng 连续生成 150 道题，这里每个长度用自己的 rng，互不影响。
    for size in (1e9, 1e10, 1e11):
        executor = ToyLLM(size, cot=False)
        print("=" * 78)
        print(f"{executor.name} executor（一次调用能可靠做约 {executor.depth_capacity()} 步）：成功率 / 平均 LLM 调用次数")
        print("=" * 78)
        print(pad("步数 k", 8) + pad("iterative executor", 22, "right") + pad("plan-and-execute", 20, "right")
              + pad("ADaPT", 16, "right") + pad("ADaPT 平均最大深度", 20, "right"))
        for k in lengths:
            rows = {"it": ([], []), "pe": ([], []), "ad": ([], [])}
            depth_used = []
            for p in problems[k]:
                truth = p.answer()
                for key, fn in (("it", lambda c: iterative_executor(executor, p, rng, c)),
                                ("pe", lambda c: plan_and_execute(executor, p, rng, c))):
                    c = Counter()
                    value, ok = fn(c)
                    rows[key][0].append(ok and value == truth)      # 成功 = 过程做对，而且最终答案对
                    rows[key][1].append(c.executor + c.planner)
                c, depths = Counter(), []
                value, ok = adapt(executor, p.start, list(p.steps), rng, c, 0, depths)
                rows["ad"][0].append(ok and value == truth)
                rows["ad"][1].append(c.executor + c.planner)
                depth_used.append(max(depths))
            cell = lambda key: f"{mean(rows[key][0]):.3f} / {mean(rows[key][1]):4.1f}"
            print(pad(str(k), 8) + pad(cell("it"), 22, "right") + pad(cell("pe"), 20, "right")
                  + pad(cell("ad"), 16, "right") + pad(f"{mean(depth_used):.2f}", 20, "right"))
        print()

    print("观察：")
    print("  · iterative executor：步数一超过 executor 的串行上限，成功率就跌到接近 0，重试也救不回来 ——")
    print("    不是运气不好，是任务对它来说\"太复杂了\"。")
    print("  · plan-and-execute 的固定粒度（每块 2 步）对短任务够用，但计划不会调整：某一块两次都失败，整个任务就失败。")
    print("    对强 executor 它又是浪费：100B 本来能一口气做 4 步，短任务上调用次数是 ADaPT 的两三倍。")
    print("  · ADaPT 按需分解：简单的任务一次做完（深度接近 0），难的任务才拆，拆到 executor 能做为止。")
    print("    同样的任务长度，executor 越弱，平均分解深度越大 —— 分解的粒度自动适应了 executor 的能力。")
    print("    多数情况下 ADaPT 的成功率明显高于固定的计划（例如 10B 做 10 步：0.36 vs 0.16；100B：0.89 vs 0.65）；")
    print("    强 executor 做中等长度的任务时两者相当。ADaPT 论文在 ALFWorld、WebShop、TextCraft 上")
    print("    报告的成功率最多分别高出 28.3%、27%、33%。")
    print("  · 代价：弱 executor 做长任务时，ADaPT 要调用很多次（不断地尝试、失败、再拆）。计算花在了需要的地方 ——")
    print("    和第 2 讲\"按难度分配 test-time compute\"是同一个思想。")
    print("  · 分解不是免费的：每次拆分都要调用 planner，也可能拆错（2%）；环境的误报（2%）也会让失败的尝试被当成成功。\n")


if __name__ == "__main__":
    main()