"""
=============================================================================
CS349E · 第 15 讲 · 02_feedback_information.py —— 反馈的信息量：工具调用一次能告诉你多少
=============================================================================

【这个脚本要回答的问题】
    KernelBench 发现，把执行结果和 profiler 输出喂回给模型，迭代优化的效果大幅提升；剩下的失败多是"结果不对"，
    因为"不对"这个反馈太粗。能不能把"反馈有多粗"量化？
      演示 1：每种（工具, 形状）在"什么都不知道"时能带来多少比特的信息，每秒多少比特。
      演示 2：几种"下一步调用哪个工具"的策略：找到 bug 的成功率、调用次数、花的时间（8 个 bug × 10 个随机种子）。
      演示 3：只说对错的测试为什么不够：哪些 bug 在它眼里一模一样（混淆）。

【关于数字】bug 的现象来自真实的 numpy 计算；工具的耗时是假设值；"策略"是固定的规则，不是语言模型。
    信息量用香农熵（比特）计算：8 个同样可能的 bug 是 3 比特；一个只说 PASS / FAIL 的测试一次最多 1 比特。
【运行】python3 02_feedback_information.py      （约 3 秒）
【依赖】numpy；本目录的 agent_lab；../lecture01 的 console。
"""

import math
import os
import sys

import numpy as np

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

from agent_lab import ACTIONS, BUG_NAMES, SHAPES, TOOLS, info_gain, likelihood_table, run_policy  # noqa: E402
from console import print_header, print_observations, print_table  # noqa: E402

SEEDS = 10


def demo1_bits(table):
    print_header("演示 1：一次工具调用能带来多少信息（从\"8 个 bug 同样可能\"的 3 比特出发）")
    prior = {b: 1 / len(BUG_NAMES) for b in BUG_NAMES}
    rows = []
    gains = {a: info_gain(prior, table, a) for a in ACTIONS}
    for a in sorted(ACTIONS, key=lambda a: -gains[a]):
        rows.append([f"{a[0]}·{a[1]}", f"{gains[a]:.2f}", f"{TOOLS[a[0]][0]:.0f} s", f"{gains[a] / TOOLS[a[0]][0]:.3f}"])
    print_table(["动作", "期望信息量（比特）", "耗时", "每秒比特"], rows)
    best = max(ACTIONS, key=lambda a: gains[a])
    t_nice, t_edge = gains[("test", "nice")], gains[("test", "edge")]
    print_observations([
        f"信息最多的是 {best[0]}·{best[1]}（{gains[best]:.2f} 比特，离完全确定只差 {3 - gains[best]:.2f} 比特）：错误的**模式**把 bug 分成了好几组。",
        f"只说对错的测试一次最多 1 比特：test·nice 只有 {t_nice:.2f} 比特，test·edge {t_edge:.2f} 比特——同样是测试，换一个形状信息量就不同，"
        "因为\"哪些 bug 在这个形状上会出错\"不同。",
        "sanitizer 类工具（memcheck、racecheck）各自只回答一个很窄的问题，信息量小、耗时长，"
        "但它们回答的问题别的工具回答不了（演示 3）。挑工具要同时看\"能告诉我多少\"和\"要等多久\"。",
    ])


def demo2_policies(table):
    print_header(f"演示 2：七种挑工具的策略（8 个 bug × {SEEDS} 个随机种子，每次最多 8 步）")
    cheap = [a for a in ACTIONS if a[0] in ("test", "test_verbose", "test_identity", "run_twice")]
    policies = [
        ("P1", "只跑 test·nice", [("test", "nice")], None),
        ("P2", "test 轮流用三种形状", [("test", s) for s in SHAPES], None),
        ("P3", "test_verbose 轮流用三种形状", [("test_verbose", s) for s in SHAPES], None),
        ("P4", "信息量最大（所有工具）", "greedy_bits", None),
        ("P5", "每秒信息量最大（所有工具）", "greedy_rate", None),
        ("P6", "每秒信息量最大（不用 sanitizer）", "greedy_rate", cheap),
        ("P7", "每秒信息量最大（不用 sanitizer，也不用构造的输入）", "greedy_rate",
         [a for a in cheap if a[0] != "test_identity"]),
    ]
    rows, res = [], {}
    for key, name, pol, allowed in policies:
        found, calls, secs, per_bug = 0, [], [], {}
        for b in BUG_NAMES:
            for seed in range(SEEDS):
                r = run_policy(pol, b, table, seed=seed, allowed=allowed)
                found += r["found"]
                per_bug[b] = per_bug.get(b, 0) + r["found"]
                calls.append(len(r["steps"]))
                secs.append(r["seconds"])
        n = len(BUG_NAMES) * SEEDS
        res[key] = (found / n, float(np.mean(calls)), float(np.mean(secs)), per_bug)
        rows.append([key, name, f"{found / n:.0%}", f"{np.mean(calls):.1f}", f"{np.mean(secs):.0f} s"])
    print_table(["", "策略", "找到 bug 的比例", "平均调用次数", "平均工具时间"], rows, aligns=["<", "<", ">", ">", ">"])
    print()
    print_table(["bug"] + [p[0] for p in policies],
                [[b] + [f"{res[p[0]][3][b]}/{SEEDS}" for p in policies] for b in BUG_NAMES])
    p4, p5, p6, p7 = res["P4"], res["P5"], res["P6"], res["P7"]
    miss7 = [b for b in BUG_NAMES if p7[3][b] < SEEDS]
    if abs(p5[2] - p4[2]) < 1:
        rate_text = (f"\"每秒信息量最大\"和\"信息量最大\"在这里几乎一样（{p5[0]:.0%}、平均 {p5[2]:.0f} s vs {p4[0]:.0%}、{p4[2]:.0f} s）："
                     "信息最多的工具恰好也是最便宜的（详细测试）。如果最有信息的工具很贵，两者就会分开——智能体的\"下一步做什么\"是一个信息 / 时间的权衡。")
    else:
        rate_text = (f"\"每秒信息量最大\"平均花 {p5[2]:.0f} s、\"信息量最大\"花 {p4[2]:.0f} s，找到的比例 {p5[0]:.0%} vs {p4[0]:.0%}："
                     "工具有成本时，智能体的\"下一步做什么\"是一个信息 / 时间的权衡。")
    print_observations([
        f"只说对错、只用一种形状，8 步也找不到 bug（{res['P1'][0]:.0%}）；轮流用三种形状好一些（{res['P2'][0]:.0%}）；"
        f"换成详细测试，{res['P3'][0]:.0%}。反馈越细，需要的调用越少。",
        rate_text,
        f"拿掉 sanitizer 仍然是 {p6[0]:.0%}：这几个 bug 靠便宜的测试就能分出来（racecheck 抓的竞争，\"跑两次结果不同\"和错误模式也能看出来）。"
        f"再拿掉\"构造的输入\"，降到 {p7[0]:.0%}，找不到的是 " + "、".join(miss7) + "：它们在随机输入上的表现一模一样（几乎全错、误差很大），"
        "只有喂一个单位矩阵——B 转置之后还是它自己——才分得开。给智能体配什么工具，决定了它**能**发现什么。",
    ])


def demo3_confusion(table):
    print_header("演示 3：只说对错时，哪些 bug 看起来一模一样")
    sig = {}
    for b in BUG_NAMES:
        key = tuple(max(table[("test", s)][b], key=table[("test", s)][b].get) for s in SHAPES)
        sig.setdefault(key, []).append(b)
    rows = [["/".join("F" if v else "p" for v in k), "、".join(v)] for k, v in sig.items()]
    print_table(["三种形状上的结果（nice/edge/bigK，F = FAIL）", "分不开的 bug"], rows, aligns=["<", "<"])
    sig_v = {}
    for b in BUG_NAMES:
        key = tuple(max(table[("test_verbose", s)][b], key=table[("test_verbose", s)][b].get) for s in SHAPES)
        sig_v.setdefault(key, []).append(b)
    groups_v = [v for v in sig_v.values() if len(v) > 1]
    print(f"\n换成详细测试（三种形状最常见的观察）：{len(sig_v)} 组" +
          ("；仍然分不开的：" + "；".join("、".join(g) for g in groups_v) if groups_v else "，每个 bug 都能分开"))
    n_groups = len(sig)
    print_observations([
        f"三种形状上的对错一共只分出 {n_groups} 组：例如 k_skip_last、fp16_accum、b_transposed 在三种形状上都错，只看对错永远分不开——"
        "可它们的修法完全不同（改循环边界 / 改累加器的类型 / 改下标）。",
        "详细测试把 fp16_accum 分了出来（零星的小误差），但 " + ("、".join(groups_v[0]) if groups_v else "") +
        " 在随机输入上仍然一样（几乎全错、误差大）；要靠\"构造的输入\"（演示 2 的 P6 vs P7）。"
        "这就是 SWE-agent 说的\"反馈要有信息量\"：一次调用告诉智能体的东西越多，它需要的轮数越少、上下文越短、花的钱越少。",
        f"信息论给了一个下界：8 个同样可能的 bug 是 log2 8 = 3 比特，一个只说对错的测试一次最多 1 比特，"
        f"所以至少要 {math.ceil(3 / 1)} 次——而且前提是每次测试都恰好把剩下的候选对半分，现实中远做不到。",
    ])


if __name__ == "__main__":
    tab = likelihood_table()
    demo1_bits(tab)
    demo2_policies(tab)
    demo3_confusion(tab)