"""
=============================================================================
第 8 讲 · 01：AlphaCode —— 大规模采样、按样例过滤、按行为聚类、只交 10 个
=============================================================================

对应讲义：lectures/lecture08_search_deep_research.md 第 2 节（Li et al. 2022；AlphaCode 2）

-----------------------------------------------------------------------------
【这个文件的小世界：猜一个隐藏的函数】
-----------------------------------------------------------------------------
每道"竞赛题"背后是一个隐藏的函数 f(x)，x 是 0~99 的整数。f 由 5~8 步运算组成，每步是下面之一：
        + c（加）、× c（乘）、mod m（取余）、min(x, c)（截断）、// d（整除）
题目给出 2 个样例（输入, 输出），样例的输入都很小（0~10），就像真实竞赛题给的样例一样。
选手要写出一个程序，在**隐藏的测试**（50 个 0~99 的随机输入）上全部正确才算解出。

模型写程序时，每一步有几种可能：
    · 写对（概率 0.30）；
    · "差一"：常数写错一点（mod m 写成 mod (m+1)、min(x, c) 写成 min(x, c−1)……）；
    · 用错运算：例如把 min 写成 max、把 mod 写成减法；
    · 漏掉这一步。
注意：很多错误在**某些输入**上看不出来。例如 min(x, 50) 和 min(x, 49) 在 x < 49 时完全一样；
x mod 37 和 x mod 38 在 x < 37 时也一样。样例的输入又很小，很多边界错误根本触发不了 —— 这正是真实竞赛题的情况
（AlphaCode 论文发现，只用题目原来的测试时，通过的程序里 62% 其实是错的）。

-----------------------------------------------------------------------------
【AlphaCode 的挑选流程】
-----------------------------------------------------------------------------
    1. 采样 k 个程序；
    2. 过滤：只保留通过全部样例的程序（AlphaCode 里这一步去掉了约 99% 的样本）；
    3. 聚类：另外生成一批**没有答案**的输入（AlphaCode 用一个专门的模型生成测试输入），
       在这些输入上运行剩下的程序，输出完全相同的程序归为一类（行为相同 = 很可能是同一种解法）；
    4. 从最大的 10 个类里各挑一个提交。10 个里有一个通过隐藏测试就算解出（"10@k"）。

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

import random

OPS = ("add", "mul", "mod", "min", "div")
P_CORRECT, P_OFF_BY_ONE, P_WRONG_OP = 0.30, 0.38, 0.19      # 剩下的 0.13：漏掉这一步
EXAMPLE_MAX = 10           # 题目给的样例都是很小的输入（0~10），像真实竞赛题的样例一样
N_HIDDEN, N_CLUSTER_INPUTS = 50, 30


# =============================================================================
# 题目与程序
# =============================================================================
def random_step(rng):
    """随机生成一步运算：(运算, 常数)。常数的范围按运算选取，让函数的值保持在合理的范围里。"""
    op = OPS[rng.randrange(len(OPS))]
    c = {"add": rng.randint(1, 30), "mul": rng.randint(2, 4), "mod": rng.randint(20, 60),
         "min": rng.randint(20, 80), "div": rng.randint(2, 5)}[op]
    return op, c


def run(program, x):
    """
    运行一个程序：程序是 (运算, 常数) 的列表，依次作用在 x 上。

    【边界】mod 和 div 的常数保证 ≥ 2（不会除以 0）；"wrong op"里 max、sub 也都安全。
    """
    for op, c in program:
        if op == "add":
            x = x + c
        elif op == "mul":
            x = x * c
        elif op == "mod":
            x = x % c
        elif op == "min":
            x = min(x, c)
        elif op == "div":
            x = x // c
        elif op == "max":
            x = max(x, c)
        elif op == "sub":
            x = x - c
    return x


WRONG_OP = {"add": "sub", "mul": "add", "mod": "sub", "min": "max", "div": "sub"}


def sample_program(problem, rng):
    """
    模拟模型写一个程序：逐步决定这一步写对、差一、用错运算还是漏掉。

    【思路】这些错误类型对应真实代码里最常见的 bug：边界差一（off-by-one）、用错函数、漏掉一个条件。
        "差一"的方向是随机的（+1 或 −1）；mod / div 的常数不会被改成小于 2 的数。
    """
    prog = []
    for op, c in problem["steps"]:
        u = rng.random()
        if u < P_CORRECT:
            prog.append((op, c))
        elif u < P_CORRECT + P_OFF_BY_ONE:
            c2 = c + (1 if rng.random() < 0.5 else -1)
            prog.append((op, max(2, c2) if op in ("mod", "div") else c2))
        elif u < P_CORRECT + P_OFF_BY_ONE + P_WRONG_OP:
            prog.append((WRONG_OP[op], c))
        # 否则：漏掉这一步
    return tuple(prog)


def make_problem(rng):
    """生成一道题：隐藏的步骤、2 个样例、50 个隐藏测试的输入。"""
    steps = [random_step(rng) for _ in range(rng.randint(5, 8))]
    examples = [rng.randint(0, EXAMPLE_MAX) for _ in range(2)]
    hidden = [rng.randint(0, 99) for _ in range(N_HIDDEN)]
    p = {"steps": steps, "examples": examples, "hidden": hidden}
    p["example_out"] = [run(steps, x) for x in examples]
    p["hidden_out"] = [run(steps, x) for x in hidden]
    return p


def passes_examples(problem, prog):
    return all(run(prog, x) == y for x, y in zip(problem["examples"], problem["example_out"]))


def passes_hidden(problem, prog):
    return all(run(prog, x) == y for x, y in zip(problem["hidden"], problem["hidden_out"]))


# =============================================================================
# 挑选方法
# =============================================================================
def select(problem, programs, method, rng, cluster_inputs):
    """
    从 programs 里挑出最多 10 个提交，返回是否有一个通过隐藏测试。

    【参数】method："random"（随机挑 10 个）/ "filter"（通过样例的里面随机挑 10 个）/
        "cluster"（通过样例的程序按行为聚类，从最大的 10 个类里各挑一个）；
        cluster_inputs：聚类用的、没有答案的输入
    【边界】通过样例的程序不足 10 个时，全部提交。
    """
    if method == "random":
        chosen = rng.sample(programs, min(10, len(programs)))
        return any(passes_hidden(problem, p) for p in chosen)
    survivors = [p for p in programs if passes_examples(problem, p)]
    if not survivors:
        return False
    if method == "filter":
        chosen = rng.sample(survivors, min(10, len(survivors)))
        return any(passes_hidden(problem, p) for p in chosen)
    clusters = {}
    for p in survivors:
        clusters.setdefault(tuple(run(p, x) for x in cluster_inputs), []).append(p)
    biggest = sorted(clusters.values(), key=len, reverse=True)[:10]
    return any(passes_hidden(problem, group[0]) for group in biggest)


def main():
    rng = random.Random(190)
    problems = [make_problem(rng) for _ in range(60)]
    ks = (10, 100, 1000, 10000)
    results = {m: {k: [] for k in ks} for m in ("oracle", "random", "filter", "cluster")}
    stats = {k: {"removed": [], "survive_correct": []} for k in ks}
    for prob in problems:
        cluster_inputs = [rng.randint(0, 99) for _ in range(N_CLUSTER_INPUTS)]
        pool = [sample_program(prob, rng) for _ in range(max(ks))]
        for k in ks:
            progs = pool[:k]
            correct = [passes_hidden(prob, p) for p in progs]
            results["oracle"][k].append(any(correct))
            for m in ("random", "filter", "cluster"):
                results[m][k].append(select(prob, progs, m, rng, cluster_inputs))
            surv = [c for p, c in zip(progs, correct) if passes_examples(prob, p)]
            stats[k]["removed"].append(1 - len(surv) / k)
            if surv:
                stats[k]["survive_correct"].append(sum(surv) / len(surv))

    mean = lambda xs: sum(xs) / len(xs)
    print("=" * 78)
    print("演示：60 道题，每题采样 k 个程序，只能提交 10 个 —— 解出的比例（10@k）")
    print("=" * 78)
    names = (("oracle", "pass@k（上限：k 个里有一个对）"), ("random", "随机提交 10 个"),
             ("filter", "按样例过滤后随机 10 个"), ("cluster", "过滤 + 聚类（AlphaCode）"))
    print(f"{'':34}" + "".join(f"k = {k}".rjust(11) for k in ks))
    for key, label in names:
        width = sum(2 if ord(ch) > 127 else 1 for ch in label)
        print(label + " " * max(1, 34 - width) + "".join(f"{mean(results[key][k]):11.3f}" for k in ks))
    print()
    print(f"{'被样例过滤掉的比例':26}" + "".join(f"{mean(stats[k]['removed']):11.3f}" for k in ks))
    print(f"{'通过样例的程序里真正正确的比例':22}" + "".join(f"{mean(stats[k]['survive_correct']):11.3f}" for k in ks))

    print("\n观察：")
    print("  · 只能交 10 个时，随机提交几乎没用：k 再大，随机挑的 10 个里也很少有对的。")
    print("  · 2 个小样例就能过滤掉约 94% 的程序 —— 但通过样例的程序里，只有大约四分之一真正正确：")
    print("    差一的常数、换错的运算，在 0~10 这样的小输入上常常看不出来。样例测试是一个有大量 false positive 的 verifier。")
    print("  · 过滤后随机挑 10 个，采样再多也提高得很慢：幸存者里错误的太多，10 个名额常常全部落在错误的程序上。")
    print("  · 聚类把\"行为相同\"的程序归为一类。正确的程序写法可能不同，在任何输入上输出都一样，会聚成一个类；")
    print("    各种错误则分散在很多不同的类里。从最大的 10 个类里各交一个，10 次机会就分给了 10 种**不同的**行为，")
    print("    而不是浪费在同一种错误的 10 个副本上。采样越多，正确的那个类越大、越容易被选中：")
    print("    k = 10000 时，过滤 + 聚类解出 0.90，过滤后随机挑只有 0.77，上限是 0.98。")
    print("  · AlphaCode 论文：解出率随采样数大致按对数线性增长；AlphaCode 2 在聚类之后再用一个打分模型")
    print("    给候选打分，在 Codeforces 上解出 43% 的题（AlphaCode 是 25%），估计在第 85 百分位。\n")


if __name__ == "__main__":
    main()