"""
=============================================================================
第 13 讲 · 01：AlphaGeometry 的配方 —— 符号推演引擎 + 从合成数据学会"辅助构造"
=============================================================================

对应讲义：lectures/lecture13_superhuman_reasoning.md 第 2 节（Trinh et al. 2024, Nature）

-----------------------------------------------------------------------------
【一个抽象的"几何"】
-----------------------------------------------------------------------------
真实的平面几何太复杂，这里用一个抽象的推理系统代替，但保留 AlphaGeometry 面对的结构：
    · 事实（fact）：题目给出的条件（编号 0~23，"基本事实"），和能推出来的结论（5 层，每层 20 个）。
    · 推演规则：形如"A 且 B ⇒ C"（像"两组对边分别平行 ⇒ 平行四边形"）。每个结论有 3 条规则可以推出它。
    · 辅助构造（auxiliary construction）：60 种（像"取 AB 的中点""过 C 作 AB 的平行线"）。每种构造会引入
      两个新的事实（新的点带来的新关系）。有些规则的前提里用到这些新事实 —— 不做构造，这些规则就永远用不上。
    · 符号推演引擎（deductive database，DD）：反复应用规则，直到推不出新的事实（"闭包"）。它是完备的，
      但**不会发明新的对象**：题目里没有的点，它不会自己加。
一道题 = 一组条件 + 一个要证的结论。有的题 DD 直接能证；有的必须先做 1 个或 2 个辅助构造。
难点：构造有 60 种、两两组合有 1770 种，而且通常看不出哪个有用 —— 这正是 IMO 几何题"灵光一现"的地方。

-----------------------------------------------------------------------------
【AlphaGeometry 的做法（toy 版）】
-----------------------------------------------------------------------------
    1. 合成数据（不用任何人类示范）：随机选一组条件 + 随机做 2 个构造，用 DD 推出所有结论；
       对每个结论**回溯**（traceback）它的证明，看证明用到了哪些前提。如果证明用到了某个构造引入的事实，
       而且不用构造就证不出来，这个构造就是"辅助构造"—— 得到一条训练数据：(结论, 需要的构造)。
    2. "语言模型"：toy 版用计数代替 —— 记住"要证结论 g 时，哪些构造曾经帮上忙"，按次数排序提出构造。
    3. 求解：先跑 DD；证不出来，就按语言模型的建议依次尝试构造（先单个、再两个一组），每次加完都跑一遍 DD。
比较的提议者：
    · 只用 DD（不做构造）；
    · 随机提构造；
    · "人设计的启发式"：按构造在规则库里出现的次数排序（"中点、平行线通常有用"）—— 和题目无关的通用知识；
    · 用合成数据训练的提议者（合成数据从 100 到 10000 个随机"图形"）。
测试题：300 道"竞赛题"（要证的结论在第 3~5 层），按同样的随机过程生成；不需要构造、需要 1 个、需要 2 个的各 100 道。

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

import random

N_BASE, N_LEVELS, PER_LEVEL, N_CONS = 24, 5, 20, 60
RULES_PER_FACT = 3          # 每个结论有几条规则能推出它
P_AUX = 0.35                # 一条规则的第二个前提是"构造引入的事实"的概率
N_PREMISES = 10             # 一道题给出几个条件
AUX_BASE = 1000             # 构造 c 引入的两个事实编号是 1000 + 2c 和 1000 + 2c + 1


def make_world(rng):
    """
    随机生成推演规则。

    【返回】(rules, level_of)
        rules   ：[(前提的 frozenset, 结论), ...]
        level_of：{事实: 层数}，基本事实是第 0 层
    【思路】
        · 第 L 层的结论由"一个第 L−1 层的事实 + 一个更低层的事实"推出（保证推理有深度）。
        · 每个结论有两个"相关的构造"；规则的第二个前提以 P_AUX 的概率换成这两个构造之一引入的事实。
          所以哪个构造有用，和要证的结论有关 —— 这是提议者能学到的规律。
        · 构造的"流行程度"不均匀（第 0 号最常用，越往后越少用，像"取中点"比"作某个奇怪的圆"常用得多），
          这是"人设计的启发式"能利用的通用知识。
    【Python 小知识】frozenset 是不可修改的集合，可以用 <= 判断"是不是子集"：{1, 2} <= {1, 2, 3} 为 True。
    """
    popularity = [1.0 / (c + 1) ** 0.8 for c in range(N_CONS)]
    total = sum(popularity)
    popularity = [p / total for p in popularity]

    def pick_construction():
        u, acc = rng.random(), 0.0
        for c, p in enumerate(popularity):
            acc += p
            if u < acc:
                return c
        return N_CONS - 1

    level_of = {b: 0 for b in range(N_BASE)}
    by_level = {0: list(range(N_BASE))}
    next_id = N_BASE
    for level in range(1, N_LEVELS + 1):
        by_level[level] = list(range(next_id, next_id + PER_LEVEL))
        for f in by_level[level]:
            level_of[f] = level
        next_id += PER_LEVEL
    rules = []
    for level in range(1, N_LEVELS + 1):
        lower = [x for lv in range(level) for x in by_level[lv]]
        for f in by_level[level]:
            relevant = [pick_construction(), pick_construction()]
            for _ in range(RULES_PER_FACT):
                a = rng.choice(by_level[level - 1])
                b = rng.choice(lower)
                if rng.random() < P_AUX:
                    c = rng.choice(relevant)
                    b = AUX_BASE + 2 * c + rng.randrange(2)
                rules.append((frozenset((a, b)), f))
    return rules, level_of


def index_rules(rules):
    """为了让 DD 跑得快：{事实: 用到它的规则下标列表}。"""
    uses = {}
    for i, (premises, _) in enumerate(rules):
        for p in premises:
            uses.setdefault(p, []).append(i)
    return uses


def deduce(facts, rules, uses):
    """
    符号推演引擎（DD）：求闭包，并记录每个结论"第一次是怎样推出来的"。

    【参数】facts：初始的事实（条件 + 构造引入的事实）
    【返回】(known, why)：known 是闭包（set）；why[结论] = 推出它的那条规则的前提
    【思路】工作队列（worklist）：每得到一个新事实，只检查**用到它**的规则，而不是每次扫描所有规则。
    【边界】推演一定会停：事实的个数有限，每个事实只进队列一次。
    """
    known, why = set(facts), {}
    queue = list(known)
    while queue:
        x = queue.pop()
        for i in uses.get(x, ()):
            premises, conclusion = rules[i]
            if conclusion not in known and premises <= known:
                known.add(conclusion)
                why[conclusion] = premises
                queue.append(conclusion)
    return known, why


def traceback(fact, why):
    """
    回溯一个结论的证明用到的**初始事实**（证明树的叶子）。

    【返回】set
    【思路】AlphaGeometry 的 traceback 会找出"证明这个结论最少需要哪些前提"；这里沿着 why 记录的那一条推导往回走。
    """
    leaves, stack = set(), [fact]
    while stack:
        x = stack.pop()
        if x in why:
            stack.extend(why[x])
        else:
            leaves.add(x)
    return leaves


def aux_facts(constructions):
    return [AUX_BASE + 2 * c + i for c in constructions for i in range(2)]


def synthetic_data(n, rules, uses, rng):
    """
    生成 n 个随机"图形"，从中提取 (结论, 辅助构造) 的训练数据。

    【返回】{结论: {构造: 次数}}
    【思路】随机条件 + 随机 2 个构造 → 推出所有结论 → 对每个结论回溯证明 →
        证明用到了构造引入的事实，而且不做构造证不出来 → 这些构造是这个结论的"辅助构造"。
        注意：整个过程没有用到任何人写的证明。
    """
    counts = {}
    for _ in range(n):
        premises = rng.sample(range(N_BASE), N_PREMISES)
        constructions = rng.sample(range(N_CONS), 2)
        known, why = deduce(premises + aux_facts(constructions), rules, uses)
        known_without, _ = deduce(premises, rules, uses)
        for f in known:
            if N_BASE <= f < AUX_BASE and f not in known_without:
                used = {(x - AUX_BASE) // 2 for x in traceback(f, why) if x >= AUX_BASE}
                table = counts.setdefault(f, {})
                for c in used:
                    table[c] = table.get(c, 0) + 1
    return counts


def make_test_set(n, rules, uses, level_of, rng):
    """
    生成 n 道"竞赛题"：要证的结论在第 3~5 层；记录它需要几个构造（0 = DD 直接能证）。

    【返回】[(条件列表, 结论), ...] 和对应的"需要的构造数"列表
        需要的构造数：0 = DD 直接能证；1 = 60 种构造里任何**一个**构造就够；2 = 单个构造都不够（生成时用的那一对可以）。
    【边界】随机选的结论可能用任何构造都证不出来（构造得太少），这样的题丢掉重来。
    """
    problems, needs = [], []
    while len(problems) < n:
        premises = rng.sample(range(N_BASE), N_PREMISES)
        constructions = rng.sample(range(N_CONS), 2)
        known, _ = deduce(premises + aux_facts(constructions), rules, uses)
        goals = sorted(f for f in known if N_BASE <= f < AUX_BASE and level_of[f] >= 3)
        if not goals:
            continue
        goal = rng.choice(goals)
        known_without, _ = deduce(premises, rules, uses)
        if goal in known_without:
            need = 0
        elif any(goal in deduce(premises + aux_facts([c]), rules, uses)[0] for c in range(N_CONS)):
            need = 1
        else:
            need = 2
        if needs.count(need) >= n // 3:
            continue                                  # 三类题各占三分之一（像精选出来的竞赛题）
        problems.append((premises, goal))
        needs.append(need)
    return problems, needs


def candidate_order(scores, rng, n_top=8):
    """
    把单个构造和两个一组的构造排成尝试的顺序（beam search 的简化版）。

    【参数】scores：{构造: 分数}（分数越高越先试）；n_top：组合时先考虑前几名
    【返回】[(c,), (c1, c2), ...]，顺序是：
        前 n_top 名的单个构造 → 它们两两组合（按分数之积排序）→ 其余的单个构造 →
        前 2·n_top 名里还没试过的两两组合。
    【思路】AlphaGeometry 用 beam search：语言模型提出几个最可能的构造，每个加上去跑一遍 DD；
        证不出来，就在这几个分支上再提下一个构造。所以"最有希望的几个构造的组合"会比
        "排名靠后的单个构造"更早被尝试。分数相同时用随机数打破平局，"随机提议者"（所有分数都相同）就是随机顺序。
    """
    tie = {c: rng.random() for c in range(N_CONS)}
    singles = sorted(range(N_CONS), key=lambda c: (-scores.get(c, 0.0), tie[c]))

    def pairs(pool):
        out = [(a, b) for i, a in enumerate(pool) for b in pool[i + 1:]]
        out.sort(key=lambda p: -(scores.get(p[0], 0.0) + 1e-3) * (scores.get(p[1], 0.0) + 1e-3))
        return out

    first = pairs(singles[:n_top])
    tried = set(first)
    later = [p for p in pairs(singles[:2 * n_top]) if p not in tried]
    return [(c,) for c in singles[:n_top]] + first + [(c,) for c in singles[n_top:]] + later


def solve(problem, proposer_scores, budget, rules, uses, rng):
    """
    先跑 DD；证不出来，就按提议者的顺序尝试最多 budget 次构造。

    【参数】proposer_scores：None 表示只用 DD；否则是 {构造: 分数}
    【返回】True / False
    """
    premises, goal = problem
    if goal in deduce(premises, rules, uses)[0]:
        return True
    if proposer_scores is None:
        return False
    for attempt in candidate_order(proposer_scores, rng)[:budget]:
        if goal in deduce(premises + aux_facts(attempt), rules, uses)[0]:
            return True
    return False


def learned_scores(counts, goal, rule_counts):
    """
    学到的提议者对每个构造的打分。

    【参数】counts：synthetic_data 的结果；goal：要证的结论；rule_counts：{构造: 在规则库里出现的次数}
    【返回】{构造: 分数}
    【思路】主要看"合成数据里，证这个结论时每个构造帮过几次忙"；再加上一个很小的通用分数
        （规则库里的出现次数 × 10⁻⁶）。合成数据里没见过的结论、或者帮过忙的构造都试完了，就退回通用的顺序。
    """
    scores = {c: v * 1e-6 for c, v in rule_counts.items()}
    for c, v in counts.get(goal, {}).items():
        scores[c] = scores.get(c, 0.0) + v
    return scores


def main():
    rng = random.Random(1300)
    rules, level_of = make_world(rng)
    uses = index_rules(rules)
    test, needs = make_test_set(300, rules, uses, level_of, random.Random(1301))
    print("=" * 78)
    print("300 道\"竞赛题\"：" + "，".join(f"需要 {k} 个构造的 {needs.count(k)} 道" for k in (0, 1, 2)))
    print("=" * 78)

    rule_counts = {}
    for premises, _ in rules:
        for p in premises:
            if p >= AUX_BASE:
                c = (p - AUX_BASE) // 2
                rule_counts[c] = rule_counts.get(c, 0) + 1
    proposers = [("只用 DD（不做构造）", None, None), ("DD + 随机提构造", {}, None),
                 ("DD + 人设计的启发式（通用的常用构造）", {c: float(v) for c, v in rule_counts.items()}, None)]
    for n_data in (100, 1000, 10000):
        proposers.append((f"DD + 学到的提议者（合成数据 {n_data}）", None, n_data))

    budgets = (1, 4, 16, 64, 256)
    print(f"{'解出的比例':40}" + "".join(f"{'尝试 ' + str(b) + ' 次':>10}" for b in budgets))
    data_rng = random.Random(1302)
    breakdown = []
    for name, scores, n_data in proposers:
        if n_data is not None:
            counts = synthetic_data(n_data, rules, uses, data_rng)
        row = ""
        for b in budgets:
            solve_rng = random.Random(1303)
            hits = []
            for problem in test:
                s = learned_scores(counts, problem[1], rule_counts) if n_data is not None else scores
                hits.append(solve(problem, s, b, rules, uses, solve_rng))
            row += f"{sum(hits) / len(hits):10.3f}"
            if b == 64:
                split = [sum(h for h, k in zip(hits, needs) if k == need) / needs.count(need) for need in (1, 2)]
                breakdown.append((name, split))
        width = sum(2 if ord(ch) > 127 else 1 for ch in name)
        print(name + " " * max(1, 40 - width) + row)

    print("\n尝试 64 次时，按题目需要的构造数拆开：")
    print(f"{'':40}{'需要 1 个':>10}{'需要 2 个':>10}")
    for name, (one, two) in breakdown[1:]:
        width = sum(2 if ord(ch) > 127 else 1 for ch in name)
        print(name + " " * max(1, 40 - width) + f"{one:10.3f}{two:10.3f}")

    print("\n观察：")
    print("  · 只用 DD 解出三分之一 —— 正好是不需要构造的那些题。DD 是完备的，但它不会发明新的对象。")
    print("    （论文：DD + AR 不做任何构造，解出 IMO-AG-30 里的 14 道。）")
    print("  · 随机提构造：只需要一个构造的题，试得够多总能碰上；需要两个构造的题，组合有 1770 种，")
    print("    随机几乎碰不上（尝试 64 次只解出 4%）。")
    print("  · 人设计的启发式只用了\"哪些构造一般比较常用\"的通用知识，已经比随机好得多。")
    print("    （论文：DD + AR 加上人设计的构造启发式，解出 18 道。）")
    print("  · 用合成数据学到的提议者最好，而且数据越多越好；提升主要在需要两个构造的题上（0.57 → 0.80）。")
    print("    它知道\"要证这一类结论时，什么构造帮过忙\"—— 这是和具体题目相关的知识，通用的启发式给不了。")
    print("    （论文：AlphaGeometry 解出 25 道，接近 IMO 金牌选手的平均 25.9 道。）")
    print("  · 整个训练过程没有用到一个人写的证明：随机\"画图\" + 符号推演 + 回溯，就足以生成教模型做构造的数据。\n")


if __name__ == "__main__":
    main()