"""
=============================================================================
第 4 讲 · 代码 06（可选）：在真实的 Claude 上比较几种思维链提示
=============================================================================

对应讲义：lectures/lecture04_chain_of_thought.md §7

【这个脚本要回答的问题】
    同一批多步算术题（2 / 4 / 8 步，各 10 道），四种提示方式：
        direct     ："Reply with only the final number."
        zero_cot   ："Let's think step by step. End with 'Answer: <number>'."
        few_cot    ：两道写好推理过程的示例 + 题目
        plan       ："First write a short plan, then carry it out step by step. End with 'Answer: <number>'."
    准确率、输出 token（真实 Claude 的输出 token 包含它在内部的"思考"）各是多少？

【一个重要的背景】默认模型 claude-opus-5-5 是一个"会思考"的模型：即使提示要求只写最终数字，它也可能先在内部推理
    （adaptive thinking；effort 控制思考的深度）。所以"direct"在这里并不等于"不推理"——推理被训练进了模型（讲义 §7）。
    可以比较 --effort low 和 --effort high 的差别。

【运行】
    python3 06_real_llm.py --toy                 # 不联网：用 reasoner_sim 的 SimReasoner-large 跑同一段代码
    python3 06_real_llm.py                       # 真实模型（effort=low）：需要 pip install anthropic 与凭证；120 次调用，预算上限 1 美元
    python3 06_real_llm.py --effort high         # 让模型想得更多（更贵）
    第二次运行会命中磁盘缓存（cache/ 目录），不再花钱。
【依赖】anthropic（可选）；同目录的 reasoner_sim.py，前置讲 P2 的 llm_api.py，第 1 讲的 console.py。
"""

import os
import random
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
for d in ("lecture01", "prereq2"):
    sys.path.insert(0, os.path.join(HERE, "..", d))
from console import pct_ci, print_header, print_observations, print_table  # noqa: E402
from llm_api import (CachedLLM, ClaudeLLM, LLMError, LLMRequest, LoggedLLM, MeteredLLM, RealClock,  # noqa: E402
                     RetryingLLM)
from reasoner_sim import SimReasoner, final_answer, make_problem  # noqa: E402

FEW_SHOT = ("Q: Start with 4. Add 3. Multiply by 2. What is the result?\nA: 4 + 3 = 7\n7 * 2 = 14\nAnswer: 14\n\n"
            "Q: Start with 9. Subtract 2. Add 5. What is the result?\nA: 9 - 2 = 7\n7 + 5 = 12\nAnswer: 12\n\n")
METHODS = {
    "direct": lambda q: (f"{q}\nReply with only the final number.", "Reply with only the final number."),
    "zero_cot": lambda q: (f"{q}\nLet's think step by step. End with 'Answer: <number>'.", ""),
    "few_cot": lambda q: (f"{FEW_SHOT}Q: {q}\nA:", "Solve the problem in the same format as the examples."),
    "plan": lambda q: (f"{q}\nFirst write a short plan, then carry it out step by step. End with 'Answer: <number>'.", ""),
}


def parse_number(text, method):
    """direct 取回答里的第一个整数；其余取 "Answer: 数字"。"""
    import re
    if method == "direct":
        m = re.search(r"-?\d+", text)
        return int(m.group(0)) if m else None
    return final_answer(text)


def main():
    toy = "--toy" in sys.argv
    effort = sys.argv[sys.argv.index("--effort") + 1] if "--effort" in sys.argv else "low"
    if toy:
        base = SimReasoner("large")
        model = base.name
    else:
        try:
            base = ClaudeLLM(effort=effort, sdk_max_retries=0)
        except LLMError as e:
            print(f"无法创建 ClaudeLLM：{e}\n可以先用 --toy 运行。")
            return
        model = base.name
    os.makedirs(os.path.join(HERE, "cache"), exist_ok=True)
    suffix = "toy" if toy else f"{model}-{effort}"
    clock = RealClock()
    retry = RetryingLLM(base, max_retries=4, base=1.0, clock=clock)
    cache = CachedLLM(retry, path=os.path.join(HERE, "cache", f"cot_{suffix}_cache.jsonl"))
    meter = MeteredLLM(cache, budget_usd=1.00)
    llm = LoggedLLM(meter, os.path.join(HERE, "cache", f"cot_{suffix}_log.jsonl"), run_id=suffix, clock=clock)

    rng = random.Random(7)
    problems = {k: [make_problem(rng, k) for _ in range(10)] for k in (2, 4, 8)}
    print_header(f"{'SimReasoner-large（--toy）' if toy else f'{model}（effort={effort}）'}：4 种提示 × 3 种步数 × 10 道题")
    rows = []
    for method, make in METHODS.items():
        row = [method]
        toks = 0
        n_all = 0
        for k, ps in problems.items():
            ok = 0
            for p in ps:
                user, system = make(p["text"])
                req = LLMRequest(messages=[{"role": "user", "content": user}], system=system, model=model,
                                 max_tokens=2000, metadata={"tag": f"{method}/{k}", "gold": p["answer"]})
                try:
                    r = llm.complete(req)
                except LLMError as e:
                    print(f"  调用失败（{type(e).__name__}）：{e}")
                    return
                ok += parse_number(r.text, method) == p["answer"]
                toks += r.usage.output_tokens
                n_all += 1
            row.append(pct_ci(ok, len(ps)))
        row.append(f"{toks / n_all:.0f}")
        rows.append(row)
    print_table(["提示方式", "2 步", "4 步", "8 步", "平均输出 token"], rows)
    print(f"\n  缓存命中 {cache.hits} 次、未命中 {cache.misses} 次；本次花费 {meter.spent:.4f} 美元")
    print_observations([
        "看 direct 一行：如果真实模型在 8 步题上 direct 也几乎全对，而输出 token 远多于'一个数字'，说明它在内部做了推理——"
        "推理模型把思维链训练进了模型，'要不要写思维链'这个提示层面的选择，变成了'让它想多深'（effort）的选择。",
        "比较各行的平均输出 token：更长的推理通常更贵、更慢；对一个已经会在内部推理的模型，再要求它在回答里写出步骤，"
        "收益可能很小，却多付了输出的钱。",
        "每格只有 10 道题，区间很宽（一道题就是 10 个百分点）；要下结论就加大题数，并用配对检验（P3）。",
        "--toy 模式下的数字由 SimReasoner 的旋钮决定，不代表任何真实模型。",
    ])


if __name__ == "__main__":
    main()