"""
=============================================================================
第 3 讲 · 代码 06（可选）：在真实的 chat 模型上测"听指令"——指令放在哪里、互相冲突时听谁的
=============================================================================

对应讲义：lectures/lecture03_instruction_tuning.md §8

【这个脚本要回答的问题】
    我们拿不到真实大模型的 base 版本，没法重复 01 号脚本的"微调前后"对比；但可以测一个指令微调过的 chat 模型
    **怎样听指令**：
    - 不给任何格式指令、在用户消息里给、在系统提示里给，回答的格式各是什么样？
    - 系统提示说"只答一个词"、用户却说"用一句完整的话解释"，模型听谁的？
    - 要求输出 JSON 时，有多少回答是合法的 JSON？

    任务和前置讲 P2 的 06 号脚本一样：把虚构国家的"资料卡"放进提示里，让模型根据资料回答（真实模型不可能知道这个世界）。
    5 种条件 × 12 道题 = 60 次调用。

【运行】
    python3 06_real_llm.py --toy     # 不联网：用 P2 的 SimLLM 跑同一段代码（它对 "only" / "one word" / "JSON" 有写死的反应）
    python3 06_real_llm.py           # 真实模型：需要 pip install anthropic，以及 ANTHROPIC_API_KEY 或 `ant auth login`
                                     # 会花钱：默认 claude-opus-5-5、effort=low，60 次短调用；预算上限 0.50 美元
    第二次运行会命中磁盘缓存（cache/ 目录），不再花钱。
【依赖】anthropic（可选）；前置讲 P2 的 llm_api.py，第 1 讲的 console.py、pretrain_world.py。
"""

import json
import os
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, SimLLM)
from pretrain_world import RELATIONS, question_text  # noqa: E402

ONE_WORD = "Answer with only the answer word."
CONDITIONS = {
    "A 没有格式指令": ("", ""),
    "B 用户消息里要求一个词": ("", ONE_WORD),
    "C 系统提示里要求一个词": (ONE_WORD, ""),
    "D 系统要一个词、用户要一句话": (ONE_WORD, "Explain your answer in a full sentence."),
    "E 系统提示要求 JSON": ('Reply only with JSON of the form {"answer": "..."}.', ""),
}


def fact_sheet(c):
    """一个国家的资料卡（英文，几句话）。和 P2 的 06 号脚本相同。"""
    return (f"Fact sheet: {c.name.capitalize()} is a country in the {c.region}. Its capital is {c.capital}. "
            f"People there speak {c.language}. It is known for its {c.export}. It is led by {c.leader}.")


def is_one_word(text):
    """回答是不是只有一个词（忽略首尾的标点和大小写）。"""
    return len(text.strip().strip(".!\"'").split()) == 1


def json_answer(text):
    """回答能不能解析成 {"answer": ...} 的 JSON；能就返回 answer 字段，否则返回 None。"""
    t = text.strip().strip("`")
    if t.startswith("json"):
        t = t[4:]
    try:
        obj = json.loads(t)
        return str(obj.get("answer", "")) if isinstance(obj, dict) else None
    except (json.JSONDecodeError, ValueError):
        return None


def main():
    toy = "--toy" in sys.argv
    sim = SimLLM()
    if toy:
        base, model = sim, sim.name
    else:
        try:
            base = ClaudeLLM(effort="low", 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 model
    clock = RealClock()
    retry = RetryingLLM(base, max_retries=4, base=1.0, clock=clock)
    cache = CachedLLM(retry, path=os.path.join(HERE, "cache", f"instr_{suffix}_cache.jsonl"))
    meter = MeteredLLM(cache, budget_usd=0.50)
    llm = LoggedLLM(meter, os.path.join(HERE, "cache", f"instr_{suffix}_log.jsonl"), run_id=suffix, clock=clock)

    print_header(f"{'SimLLM（--toy）' if toy else model}：5 种指令条件 × 12 道题")
    probes = [(c, rel) for c in sim.world.countries[:3] for rel in RELATIONS]
    rows, examples = [], []
    for cname, (system, user_suffix) in CONDITIONS.items():
        ok = one = valid = n = 0
        out_tokens = 0
        for c, rel in probes:
            q = question_text(c.name, rel).replace(" ?", "?")
            content = f"{fact_sheet(c)}\n\nQuestion: {q}" + (f"\n{user_suffix}" if user_suffix else "")
            req = LLMRequest(messages=[{"role": "user", "content": content}], system=system, model=model,
                             max_tokens=300, metadata={"tag": cname, "gold": c.attr(rel)})
            try:
                r = llm.complete(req)
            except LLMError as e:
                print(f"  调用失败（{type(e).__name__}）：{e}")
                return
            gold = c.attr(rel)
            ja = json_answer(r.text)
            valid += ja is not None
            hit = (ja or r.text).lower().replace(".", " ").replace('"', " ").split()
            ok += gold in hit
            one += is_one_word(r.text)
            out_tokens += r.usage.output_tokens
            n += 1
            if rel == "capital" and c is sim.world.countries[0]:
                examples.append([cname, repr(r.text[:70])])
        rows.append([cname, pct_ci(ok, n), pct_ci(one, n), f"{valid}/{n}", f"{out_tokens / n:.1f}"])
    print_table(["条件", "含正确答案", "只有一个词", "合法 JSON", "平均输出 token"], rows)
    print()
    print_table(["条件", "问首都时的回答（前 70 个字符）"], examples, aligns=["<", "<"])
    print(f"\n  缓存命中 {cache.hits} 次、未命中 {cache.misses} 次；本次花费 {meter.spent:.4f} 美元")
    print_observations([
        "A → B → C：同一条'只答一个词'的指令，放在用户消息里和放在系统提示里，效果是否一样？指令微调过的模型通常两种都听，"
        "但系统提示更适合放'每次都适用'的规则（P2 §3.3）。",
        "D 是一个冲突：系统提示和用户消息要求相反的格式。你的模型听了谁？很多 chat 模型被训练成优先听系统提示（指令层级），"
        "但并不保证——**不要设计让指令互相冲突的系统**。",
        "E：要求 JSON 时，'合法 JSON'一列是格式遵从率。生产系统里应该总是用代码校验（P1 §8 的 extract_json），"
        "而不是假设模型一定照做；有的 API 提供结构化输出功能，从解码层面保证格式。",
        "--toy 模式下的行为是 SimLLM 写死的规则（看到 one word 就只写一个词），不代表任何真实模型。",
    ])


if __name__ == "__main__":
    main()