"""
=============================================================================
前置讲 P2 · 代码 05：一次提示实验的"基础设施" —— 记账、缓存、日志、预算
=============================================================================

对应讲义：lectures/prereq2_llm_as_a_service.md §6

【这个脚本要回答的问题】
    一次典型的提示实验：3 个提示模板 × 96 道题 = 288 次调用。
    1. 怎样知道每个模板花了多少钱、用了多少 token？（演示 1：记账）
    2. 改了一个模板之后重跑，怎样只为改动的部分付钱？（演示 2：缓存）
    3. 一周之后，怎样证明表格里的数字是怎么来的？（演示 3：只凭日志重算结果）
    4. 怎样防止一次实验花超预算？中途停下之后怎样接着跑？（演示 4：预算）
    5. 温度 > 0 时，缓存会悄悄把实验变成什么样？（演示 5）

【怎么读输出】
    模型是 SimLLM（按虚构的 sim-large 价格计费：每百万 token 输入 3 美元、输出 15 美元），
    底下套了 FlakyLLM（10% 的 429、3% 的 500）和 RetryingLLM。时间用虚拟时钟，所以每次运行结果都一样。

【运行】python3 05_experiment_infra.py     （约 2 秒）
【依赖】只用 Python 标准库；import 同目录的 llm_api.py，前置讲 P1 的 text_lab.py，第 1 讲的 console.py、pretrain_world.py。
"""

import json
import os
import re
import sys
import tempfile

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "lecture01"))
from console import pct_ci, print_header, print_observations, print_table  # noqa: E402
from llm_api import (BudgetExceeded, CachedLLM, FlakyLLM, LLMRequest, LoggedLLM, MeteredLLM,  # noqa: E402
                     RetryingLLM, SimLLM, VirtualClock, read_log)
from pretrain_world import RELATIONS, question_text  # noqa: E402
from text_lab import extract_json  # noqa: E402

TEMPLATES = {
    "A 直接问": "",
    "B 只答一个词": "Answer with only the answer word.",
    "C 只输出 JSON": "Reply only in JSON with a single key \"answer\".",
}


def probes(world):
    """96 道题：(国家, 关系, 问题文字, 标准答案)。"""
    return [(c.name, rel, question_text(c.name, rel).replace(" ?", "?"), c.attr(rel))
            for c in world.countries for rel in RELATIONS]


def score(template, text, gold):
    """
    判断一个回答对不对。

    【思路】C 模板要求 JSON：取出 JSON 里的 answer 字段比较（取不出来就算错——格式错误也是错误）；
        A、B 模板：把回答转成小写、去掉标点，看标准答案是不是其中一个词。
    【边界】这种"词出现就算对"的判法很宽松（"The answer is not grorken." 也会被判对）。
        评估指标的设计会直接影响结论，第 17 讲专门讨论。
    """
    if template.startswith("C"):
        obj, _ = extract_json(text)
        return isinstance(obj, dict) and obj.get("answer") == gold
    return gold in re.sub(r"[^\w\s]", " ", text.lower()).split()


def build_stack(clock, cache_path=None, log_path=None, budget=None, seed=0):
    """
    组装一个完整的调用栈（从外到内）：日志 → 记账与预算 → 缓存 → 重试 → 注入失败 → 模拟模型。

    【为什么是这个顺序】
        缓存在重试**外面**：命中缓存就根本不用调用、也不会失败；
        记账在缓存外面：能看到哪些调用命中了缓存（命中的不计费）；
        日志在最外面：每一次调用（包括命中缓存的、失败的）都留下记录。
    """
    llm = SimLLM(name="sim-large")
    llm = FlakyLLM(llm, p_429=0.10, p_500=0.03, retry_after=1.0, seed=seed, clock=clock)
    llm = RetryingLLM(llm, max_retries=5, clock=clock, seed=seed)
    cache = CachedLLM(llm, path=cache_path)
    meter = MeteredLLM(cache, budget_usd=budget)
    top = LoggedLLM(meter, log_path, run_id="exp", clock=clock) if log_path else meter
    return top, meter, cache


def run(stack, world, templates=TEMPLATES, stop_on_budget=False):
    """
    跑一遍实验：每个模板 × 每道题调用一次。

    【返回】{模板名: (答对数, 题数)}；遇到 BudgetExceeded 且 stop_on_budget=True 时提前返回已完成的部分
    """
    res = {t: [0, 0] for t in templates}
    for t, system in templates.items():
        for country, rel, q, gold in probes(world):
            req = LLMRequest(messages=[{"role": "user", "content": q}], system=system, model="sim-large",
                             max_tokens=40, metadata={"tag": t, "country": country, "relation": rel, "gold": gold})
            try:
                r = stack.complete(req)
            except BudgetExceeded:
                if stop_on_budget:
                    return res
                raise
            res[t][0] += score(t, r.text, gold)
            res[t][1] += 1
    return res


def demo1_metering(world, d):
    """演示 1：第一次运行与记账。"""
    print_header("演示 1：跑一次实验——每个模板花了多少")
    clock = VirtualClock()
    stack, meter, cache = build_stack(clock, os.path.join(d, "cache.jsonl"), os.path.join(d, "log.jsonl"))
    res = run(stack, world)
    rows = []
    for t in TEMPLATES:
        n, tin, tout, cost = meter.by_tag[t]
        rows.append([t, pct_ci(*res[t]), f"{tin / n:.1f}", f"{tout / n:.1f}", f"{cost:.5f}"])
    retrier = cache.inner
    print_table(["模板", "准确率 [95% CI]", "平均输入 token", "平均输出 token", "费用（美元）"], rows)
    print(f"  合计 {meter.calls} 次调用，{meter.spent:.5f} 美元；中途重试 {retrier.retries} 次，"
          f"重试等待 {retrier.waited:.1f} 秒（虚拟时钟），总耗时 {clock.now():.1f} 秒")
    per = {t: (meter.by_tag[t][1] / meter.by_tag[t][0], meter.by_tag[t][2] / meter.by_tag[t][0], meter.by_tag[t][3])
           for t in TEMPLATES}
    a_, b_, c_ = per["A 直接问"], per["B 只答一个词"], per["C 只输出 JSON"]
    print_observations([
        f"按 metadata 里的 tag 分组记账，每个模板的成本一目了然。B 的系统提示让每次的输入多了约 {b_[0] - a_[0]:.0f} 个 token，"
        f"但输出从 {a_[1]:.1f} 降到 {b_[1]:.1f} 个 token；输出的单价是输入的 5 倍，所以 B 总体反而"
        + ("更便宜" if b_[2] < a_[2] else "更贵") + f"（{b_[2]:.5f} vs {a_[2]:.5f} 美元）。"
        "C 要求 JSON，引号和括号让输出最长，也最贵。",
        "准确率的置信区间宽达 ±7 个百分点左右：96 道题只够分辨很大的差别（前置讲 P3）。",
        f"中途有 {retrier.retries} 次重试，但调用方完全没感觉到——失败被底层的 RetryingLLM 吸收了。",
    ])
    return meter.spent, res


def demo2_cache(world, d, first_cost):
    """演示 2：重跑与缓存。"""
    print_header("演示 2：重跑实验——缓存让'没改的部分'免费")
    rows = []
    clock = VirtualClock()
    stack, meter, cache = build_stack(clock, os.path.join(d, "cache.jsonl"), os.path.join(d, "log2.jsonl"))
    run(stack, world)
    rows.append(["原样重跑一遍", cache.hits, cache.misses, f"{meter.spent:.5f}", f"{clock.now():.1f} 秒"])
    changed = dict(TEMPLATES)
    changed["B 只答一个词"] = "Answer with only the answer word, nothing else."
    clock = VirtualClock()
    stack, meter, cache = build_stack(clock, os.path.join(d, "cache.jsonl"), os.path.join(d, "log3.jsonl"))
    run(stack, world, changed)
    rows.append(["只改了模板 B 的一句话", cache.hits, cache.misses, f"{meter.spent:.5f}", f"{clock.now():.1f} 秒"])
    print(f"第一次运行花了 {first_cost:.5f} 美元。缓存存在磁盘上（JSONL），重启程序也还在。")
    print_table(["这次运行", "缓存命中", "缓存未命中", "费用（美元）", "总耗时"], rows)
    print_observations([
        "原样重跑：288 次全部命中缓存，0 费用、0 等待。调试分析代码、重画图表时，不用再付一次钱。",
        "改了模板 B 的一个字，B 的 96 个请求的指纹全变了（系统提示是请求的一部分），只有它们需要重新调用；"
        "A 和 C 照旧命中缓存。",
        "缓存命中的前提是请求**逐字节相同**：在提示里放了当前时间、随机 ID、没排序的 JSON，缓存就永远不会命中。"
        "服务端的'提示缓存'（prompt caching）也是同样的道理：它按相同的**前缀**来匹配。",
    ])


def demo3_log(world, d, res):
    """演示 3：只凭日志重算结果。"""
    print_header("演示 3：一周之后——只凭日志重算表格里的数字")
    recs = read_log(os.path.join(d, "log.jsonl"))
    again = {t: [0, 0] for t in TEMPLATES}
    for r in recs:
        t = r["metadata"]["tag"]
        again[t][0] += score(t, r["response"]["text"], r["metadata"]["gold"])
        again[t][1] += 1
    rows = [[t, pct_ci(*res[t]), pct_ci(*again[t]), "✓" if res[t] == again[t] else "✗"] for t in TEMPLATES]
    print_table(["模板", "实验时算出的准确率", "只用日志重算", "一致"], rows)
    sample = dict(recs[100])
    sample["response"] = {k: sample["response"][k] for k in ("text", "stop_reason", "usage", "cached", "attempts")}
    print("\n  日志里的一条记录：")
    print("    " + json.dumps(sample, ensure_ascii=False, indent=2).replace("\n", "\n    "))
    print_observations([
        "日志里有重算结果需要的一切：发给模型的完整请求（模型、系统提示、消息、参数）、模型的原始回答、"
        "标准答案（放在 metadata 里）。换一种判分方法（例如更严格的精确匹配），也只需要重新读日志，不用重新调用模型。",
        "attempts 记录了这次调用试了几次：如果某个模板的失败重试特别多，日志能帮你发现。",
        "实验报告里的每一个数字，都应该能从日志重新算出来。这是第 17 讲'可复现的提示实验'的第一条要求。",
    ])


def demo4_budget(world, d, first_cost):
    """演示 4：预算与续跑。"""
    print_header("演示 4：预算上限——花到一半停下，再接着跑")
    path = os.path.join(d, "cache_budget.jsonl")
    budget = round(first_cost * 0.5, 5)
    clock = VirtualClock()
    stack, meter, cache = build_stack(clock, path, budget=budget)
    res = run(stack, world, stop_on_budget=True)
    done = sum(n for _, n in res.values())
    rows = [[f"预算 {budget:.5f} 美元（第一次运行费用的一半）", done, f"{meter.spent:.5f}", "BudgetExceeded，停下"]]
    stack, meter2, cache2 = build_stack(VirtualClock(), path, budget=first_cost * 2)
    res2 = run(stack, world)
    rows.append([f"提高预算后重新运行（同一个缓存）", sum(n for _, n in res2.values()), f"{meter2.spent:.5f}",
                 f"缓存命中 {cache2.hits} 次"])
    print_table(["这次运行", "完成的调用", "花费（美元）", "说明"], rows)
    print(f"  两次合计 {meter.spent + meter2.spent:.5f} 美元；一次跑完是 {first_cost:.5f} 美元")
    print_observations([
        "MeteredLLM 在**每次调用之前**按最坏情况估算（输入 token + max_tokens 个输出 token），"
        "估计会超预算就拒绝调用。实际花费一定不超过预算。",
        "停下之后，已经完成的调用都在缓存里。提高预算重跑，前面的部分直接命中缓存，只为剩下的部分付钱："
        "两次合计正好等于一次跑完的费用。**缓存让实验可以随时中断、随时续跑。**",
        "真实实验里还要给每个请求设 max_tokens 上限、给每分钟设请求数上限：一个写错的循环可以在几分钟里花掉几百美元。",
    ])


def demo5_sampling(world):
    """演示 5：温度 > 0 时的缓存陷阱。"""
    print_header("演示 5：温度 > 0 时，缓存会悄悄'冻住'样本")
    qs = [question_text(c.name, "capital").replace(" ?", "?") for c in world.countries]
    rows = []
    for name, cache_sampled, seeds in [("缓存采样结果，不给 seed（错误）", True, [None] * 5),
                                       ("不缓存采样结果，不给 seed（默认）", False, [None] * 5),
                                       ("缓存采样结果，seed = 0–4", False, [0, 1, 2, 3, 4])]:
        cache = CachedLLM(SimLLM(name="sim-large"), cache_sampled=cache_sampled)
        distinct = 0
        for q in qs:
            outs = {cache.complete(LLMRequest(messages=[{"role": "user", "content": q}], model="sim-large",
                                              temperature=1.5, seed=s, max_tokens=30)).text for s in seeds}
            distinct += len(outs)
        before = cache.hits
        for q in qs:
            for s in seeds:
                cache.complete(LLMRequest(messages=[{"role": "user", "content": q}], model="sim-large",
                                          temperature=1.5, seed=s, max_tokens=30))
        rows.append([name, f"{distinct / len(qs):.2f}", cache.hits - before, cache.bypass])
    print(f"{len(qs)} 道题，每道题在温度 1.5 下采样 5 次；然后把同样的 5 × {len(qs)} 次采样再请求一遍")
    print_table(["做法", "每道题平均有几种不同的回答", "第二遍命中缓存的次数", "绕过缓存的次数"], rows)
    print_observations([
        "如果把'没给 seed 的采样请求'也缓存起来，5 次'采样'其实是同一个请求的 5 次命中——每道题只有 1 种回答。"
        "你以为在测模型的随机性（第 4 讲的 self-consistency 要多次采样再投票），其实在看同一个样本。",
        "CachedLLM 默认让这类请求绕过缓存：样本是真的，但每次都要付钱，也没法复现。",
        "最好的做法是给每次采样一个 seed（第几次采样就用第几号 seed）：样本各不相同，而且 seed 是请求指纹的一部分，"
        "可以缓存、可以复现。真实 API 不支持 seed 时，就把每次采样的结果连同'第几次'一起写进日志。",
    ])


def main():
    world = SimLLM().world
    with tempfile.TemporaryDirectory() as d:
        first_cost, res = demo1_metering(world, d)
        demo2_cache(world, d, first_cost)
        demo3_log(world, d, res)
        demo4_budget(world, d, first_cost)
    demo5_sampling(world)


if __name__ == "__main__":
    main()