"""
=============================================================================
第 6 讲 · 01：STaR —— 用自己答对的推理训练自己
=============================================================================

对应讲义：lectures/lecture06_scaling_rl.md 第 2 节（Zelikman et al. 2022）

-----------------------------------------------------------------------------
【STaR 的循环】
-----------------------------------------------------------------------------
    重复若干轮：
        1. 用当前模型对训练集里的每道题生成一份推理（rationale）和答案；
        2. 只保留答案**正确**的那些（需要一个 verifier：这里是核对标准答案）；
        3. rationalization（可选）：对答错的题，把正确答案作为**提示**给模型，让它"倒推"出一份推理；
           如果这份推理得到了正确答案，也保留下来（训练时不带提示 —— 假装是模型自己想出来的）；
        4. 从**原始模型**出发，在保留下来的数据上微调，得到新的模型。
    模型从自己的成功中学习，不需要人写推理过程。

-----------------------------------------------------------------------------
【三个版本】
-----------------------------------------------------------------------------
    STaR                    ：只用自己答对的推理。
    STaR + rationalization  ：再加上"看了答案提示后"答对的推理。提示的效果：每一步做对的 logit + 1.5。
    STaR + 粗糙的 verifier   ：verifier 只核对答案的**个位数**（像选择题那样，蒙对的机会是 10%）。
                              错误的推理也会混进训练数据。

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

import random

from trainable_llm import (TrainableLLM, context_report, evaluate, make_problem_set, mean, pad, sft)

HINT_BONUS = 1.5
ITERATIONS = 6


class Hinted(TrainableLLM):
    """
    "看了答案提示"的模型：在每一步 correct 的 logit 上加 HINT_BONUS。

    【思路】知道最终答案，模型更容易写出一份通向它的推理（可以倒推、可以检查）。这里用"每一步更容易做对"
        来卡通化这一点。注意：训练数据里**不带**提示，所以模型学到的是"不看答案也这样推理"。
    """

    def step_logits(self, op, a, b, step_index):
        z = super().step_logits(op, a, b, step_index)
        return [z[0] + HINT_BONUS, z[1], z[2]]


def verify(problem, sample, coarse):
    """verifier：精确核对答案，或者只核对个位数（coarse=True）。"""
    if coarse:
        return abs(sample.answer) % 10 == abs(problem.answer()) % 10
    return sample.correct


def star(base_size, train, test, rng, rationalize=False, coarse=False):
    """
    跑 ITERATIONS 轮 STaR，返回每一轮的记录。

    【返回】list of dict：{"iter", "test_pass1", "n_data", "n_wrong_kept", "contexts"}
        n_wrong_kept：训练数据里其实答错了（被粗糙的 verifier 误放进来）的解答数。
    """
    model = TrainableLLM(base_size)
    log = [{"iter": 0, "test_pass1": evaluate(model, test, rng)[1], "n_data": 0, "n_wrong_kept": 0,
            "contexts": context_report(model)}]
    for it in range(1, ITERATIONS + 1):
        data = []
        for p in train:
            s = model.generate(p, rng)
            if verify(p, s, coarse):
                data.append((p, s))
            elif rationalize:
                s2 = Hinted(model.size, model.delta).generate(p, rng)
                if verify(p, s2, coarse):
                    data.append((p, s2))
        model = sft(TrainableLLM(base_size), data)          # 每一轮都从原始模型出发微调（STaR 论文的做法）
        log.append({"iter": it, "test_pass1": evaluate(model, test, rng)[1], "n_data": len(data),
                    "n_wrong_kept": sum(1 for p, s in data if not s.correct), "contexts": context_report(model)})
    return log


def main():
    train = make_problem_set(random.Random(140), n_problems=300, min_steps=3, max_steps=8)
    test = make_problem_set(random.Random(141), n_problems=200, min_steps=3, max_steps=8)
    runs = (("STaR", dict()), ("STaR + rationalization", dict(rationalize=True)),
            ("STaR + 粗糙的 verifier", dict(coarse=True)))
    logs = {name: star(1e9, train, test, random.Random(142), **kw) for name, kw in runs}

    print("=" * 78)
    print("演示 1：STaR 的迭代 —— 测试集 pass@1（1B 模型，temperature = 1，300 道训练题、200 道测试题）")
    print("=" * 78)
    print(pad("轮次", 6) + "".join(pad(name, 26, "right") for name, _ in runs))
    for i in range(ITERATIONS + 1):
        row = pad(str(i), 6)
        for name, _ in runs:
            rec = logs[name][i]
            extra = f"（{rec['n_data']} 条" + (f"，{rec['n_wrong_kept']} 条错" if rec["n_wrong_kept"] else "") + "）" \
                if i else ""
            row += pad(f"{rec['test_pass1']:.3f}{extra}", 26, "right")
        print(row)
    print("（括号里是这一轮保留下来的训练数据条数，以及其中其实答错了的条数）")

    print("\n最后一轮，几个情境下选 correct 的概率（训练前 → 训练后）：")
    before = logs["STaR"][0]["contexts"]
    print(pad("情境", 26) + pad("训练前", 10, "right") + "".join(pad(name, 26, "right") for name, _ in runs))
    for ctx in before:
        print(pad(ctx, 26) + pad(f"{before[ctx]:.3f}", 10, "right")
              + "".join(pad(f"{logs[name][-1]['contexts'][ctx]:.3f}", 26, "right") for name, _ in runs))

    print("\n观察：")
    print("  · STaR 只用\"自己答对的推理\"作为训练数据，测试集 pass@1 在前三轮里从 0.107 涨到 0.43 —— 模型在自我改进；")
    print("    保留下来的数据也越来越多：模型越强，答对的题越多，下一轮的训练数据就越多（一个正反馈的飞轮）。")
    print("  · 但第三轮以后就停在 0.4 左右（还有一点随机波动）：总有一批难题模型始终做不对，")
    print("    它们永远产生不了训练数据，飞轮停转了。最难的情境（乘两位数、被乘数 ≥ 100）只从 0.10 升到 0.20。")
    print("  · rationalization 把\"答案提示\"变成了难题上的训练数据：数据量翻倍，最难情境的正确率从 0.10 升到 0.56，")
    print("    pass@1 也更高。STaR 论文在 n 位数加法上也发现，加上 rationalization 学得更快。")
    print("  · 粗糙的 verifier（只核对个位数）让**超过一半**的保留数据其实是错的推理。模型被这些数据\"教坏\"了：")
    print("    pass@1 只涨到 0.2 出头，最难的情境几乎原地不动 —— verifier 的 false positive 会在自我训练中")
    print("    被模型学进去（第 3 讲）。STaR 论文也讨论过这个问题：在 CommonsenseQA 这样的选择题上，蒙对的推理会混进数据。\n")


if __name__ == "__main__":
    main()