"""
=============================================================================
CS525 · 第 9 讲模块：一个可验证的推理世界（reasoning_world.py）
=============================================================================

对应讲义：lectures/lecture09_posttraining_data_ii.md

-----------------------------------------------------------------------------
【为什么要造这个世界】
-----------------------------------------------------------------------------
推理模型的训练数据有一个共同的特点：**答案可以被检查**——数学题有标准答案，代码有单元测试。
这让两件事成为可能：(1) 用"答对了没有"直接当强化学习的奖励（RLVR）；(2) 从强模型那里收集解题过程（推理轨迹），
按"最后答对了没有"筛选，再拿去微调小模型（蒸馏）。这个模块造一个答案可以被检查的小世界：

    题目     从一个数出发，依次做 k 步运算（模 101），问最后的结果。k 是难度（1–12 步）。
             运算有 6 种（加、减、乘、平方、取反、加倍），一道题里每一步的运算是随机的。
    模型     对每种运算有一个"做对的概率" p[运算]；还有一个"能连续推多少步"的视野 H：
             第 i 步如果超过了 H，额外乘一个衰减（越往后越容易乱）。一步做错，后面的数就全错了。
             所以一道 k 步的题答对的概率 = 每一步做对的概率的乘积（再加上极小的"蒙对"概率 1/101）。
    推理轨迹 老师解一道题时写下的过程：每一步的中间结果。"完整"的老师每步都写；"简略"的老师跳过一半的步骤。
    学习     蒸馏（SFT）：学生看到的每一步，把对应运算的能力往"这一步写得对不对"挪一点；
             见到比自己的视野更长的推理链，视野就往长了拉一点。
             强化学习（GRPO）：见 grpo_step。

这个世界没有"语言"，只有"每一步做没做对"——所以它适合研究数据的难度、数量、过滤、老师的选择，
不适合研究推理的格式和风格。

【运行自检】
    python3 reasoning_world.py        # 打印一个学生和一个老师在各难度上的准确率
"""

import numpy as np

MOD = 101
OPS = ["加", "减", "乘", "平方", "取反", "加倍"]
N_OPS = len(OPS)


def apply_op(op, x, c):
    """做一步运算（模 101）。op 是运算编号，c 是这一步用的常数。"""
    if op == 0:
        return (x + c) % MOD
    if op == 1:
        return (x - c) % MOD
    if op == 2:
        return (x * c) % MOD
    if op == 3:
        return (x * x) % MOD
    if op == 4:
        return (-x) % MOD
    return (2 * x) % MOD


class Problems:
    """
    一批题目：每道题 = 起始数 + k 步（运算, 常数）+ 正确答案。

    【属性】start（n,）、ops（n × 12，用不到的步填 −1）、consts、k（每道题的步数）、answer
    """

    def __init__(self, start, ops, consts, k):
        self.start, self.ops, self.consts, self.k = start, ops, consts, k
        ans = start.copy()
        for i in range(ops.shape[1]):
            m = ops[:, i] >= 0
            for j in np.flatnonzero(m):
                ans[j] = apply_op(ops[j, i], ans[j], consts[j, i])
        self.answer = ans

    def __len__(self):
        return len(self.k)

    def subset(self, idx):
        idx = np.asarray(idx)
        return Problems(self.start[idx], self.ops[idx], self.consts[idx], self.k[idx])

    def op_counts(self):
        """每道题里每种运算出现了几次（n × 6）。"""
        out = np.zeros((len(self), N_OPS), dtype=int)
        for t in range(N_OPS):
            out[:, t] = (self.ops == t).sum(1)
        return out


def make_problems(n, rng, k_range=(1, 12), op_probs=None):
    """
    生成 n 道题。

    【参数】n；rng；k_range：步数的范围（含两端，均匀抽）；op_probs：6 种运算的抽取概率（None = 均匀）
    【返回】Problems
    """
    k = rng.integers(k_range[0], k_range[1] + 1, size=n)
    ops = np.full((n, 12), -1)
    consts = np.zeros((n, 12), dtype=int)
    for j in range(n):
        ops[j, :k[j]] = rng.choice(N_OPS, size=k[j], p=op_probs)
        consts[j, :k[j]] = rng.integers(2, 10, size=k[j])
    return Problems(rng.integers(0, MOD, size=n), ops, consts, k)


class Model:
    """
    一个会解题的模型：每种运算的能力 p（6 个数）+ 视野 H + 视野之外每步的衰减 decay。

    【参数】p：6 种运算做对的概率；H：能稳稳推下去的步数；decay：超过 H 的每一步，做对的概率再乘以 decay^(超出的步数)
    """

    def __init__(self, p, H=4.0, decay=0.85):
        self.p = np.array(p, dtype=float)
        self.H = float(H)
        self.decay = decay

    def copy(self):
        return Model(self.p.copy(), self.H, self.decay)

    def step_probs(self, problems):
        """每道题每一步做对的概率（n × 12；用不到的步为 1）。"""
        P = np.ones(problems.ops.shape)
        for i in range(problems.ops.shape[1]):
            m = problems.ops[:, i] >= 0
            over = max(0.0, i + 1 - self.H)
            P[m, i] = self.p[problems.ops[m, i]] * self.decay ** over
        return P

    def accuracy(self, problems):
        """每道题答对的概率（精确计算，不用抽样）：每一步都对，或者错了但碰巧蒙对（1/101）。"""
        allright = self.step_probs(problems).prod(1)
        return allright + (1 - allright) / MOD

    def solve(self, problems, rng, show_frac=1.0):
        """
        真的解一遍（抽样），写下推理轨迹。

        【参数】problems；rng；show_frac：轨迹里写出每一步的概率（1 = 完整，0.5 = 简略的老师跳过一半步骤）
        【返回】字典：correct（每道题最后答对没有）、step_ok（n × 12，每一步的"做法"对不对）、shown（n × 12，这一步写出来了没有）
        【思路】每一步的做法独立地以 p 的概率做对。一步做错之后，后面的数都错了，最后的答案也错了——
                但后面那些步骤的**做法**仍然可能是对的（对一个错的数正确地做了运算）。所以一条答错的轨迹里，
                大部分步骤仍然在示范正确的做法：这是"按最后答案过滤会丢掉有用数据"的根源（代码 02 演示 2）。
        """
        P = self.step_probs(problems)
        used = problems.ops >= 0
        step_ok = (rng.random(P.shape) < P) | ~used
        correct = step_ok.all(1)
        lucky = (~correct) & (rng.random(len(problems)) < 1 / MOD)
        shown = used & (rng.random(P.shape) < show_frac)
        return {"correct": correct | lucky, "step_ok": step_ok & used, "shown": shown}


def distill(student, problems, traces, rate=1.0, kappa=300.0, h_kappa=200.0):
    """
    蒸馏（SFT）：学生在老师的推理轨迹上微调。

    【参数】student：Model（会被修改）；problems、traces：老师解的题和写下的轨迹（可以先按条件筛选过）；
            rate：最多能往老师的水平挪多少；kappa：每种运算要看多少步才学得明显；h_kappa：要看多少条长轨迹，视野才明显变长
    【思路】(1) 每种运算：学生看到的写出来的步骤里，这种运算有 n 步、其中做法正确的比例是 q；
                p ← p + rate · (1 − e^{−n/kappa}) · (q − p)。看不到的步骤（简略的老师跳过的）什么也学不到。
            (2) 视野：比学生视野长的轨迹（按写出来的步数算）有 m 条、平均长度 L；
                H ← H + (1 − e^{−m/h_kappa}) · (L − H)。长的推理链——哪怕最后答错了——也在示范"怎样一直推下去"。
    【返回】student
    """
    shown = traces["shown"]
    ok = traces["step_ok"] & shown
    for t in range(N_OPS):
        mask = shown & (problems.ops == t)
        n = mask.sum()
        if n == 0:
            continue
        q = ok[mask].mean()
        student.p[t] += rate * (1 - np.exp(-n / kappa)) * (q - student.p[t])
    lengths = shown.sum(1)
    long = lengths > student.H
    if long.any():
        L = lengths[long].mean()
        student.H += (1 - np.exp(-long.sum() / h_kappa)) * (L - student.H)
    return student


if __name__ == "__main__":
    rng = np.random.default_rng(0)
    test = make_problems(3000, rng)
    student = Model([0.80, 0.80, 0.70, 0.65, 0.85, 0.80], H=4)
    teacher = Model([0.97] * 6, H=12)
    for name, m in [("学生", student), ("老师", teacher)]:
        acc = m.accuracy(test)
        print(name, "平均准确率", round(float(acc.mean()), 3), "按难度：",
              {int(k): round(float(acc[test.k == k].mean()), 2) for k in (1, 3, 6, 9, 12)})
    tr = make_problems(2000, rng, (4, 12))
    distill(student, tr, teacher.solve(tr, rng))
    print("蒸馏 2000 条完整轨迹之后的学生：", round(float(student.accuracy(test).mean()), 3), "视野", round(student.H, 2))