"""
=============================================================================
第 6 讲 · 工具模块：一个可以训练的 toy LLM
=============================================================================

对应讲义：lectures/lecture06_scaling_rl.md

-----------------------------------------------------------------------------
【思路】
-----------------------------------------------------------------------------
第 1 讲的 ToyLLM 每一步按 softmax(logits) 在 correct / bug / careless 三种做法里选一种，logits 是固定的公式。
这里给它加上**可以学习的参数**：对每一种"情境"（context）—— 运算类型 × 这一步的难度（hardness）——
有一个长度为 3 的偏移量 Δ，加到原来的 logits 上：

        logits(情境) = ToyLLM 原来的 logits + Δ[情境]

    · Δ 初始为 0：训练前的模型和第 1 讲的 ToyLLM 完全一样。
    · 情境由（运算, hardness）决定，例如"乘两位数、被乘数 ≥ 100"就是 ("mul", 2.2)。
      所有题目里同一种情境**共享**同一组参数 —— 在一道题上学到的东西，会用到所有有这种情境的题上（泛化）。
    · 训练 = 调整 Δ，让"做对"的做法概率变大。这就是一个只有十几个"神经元"的语言模型。

梯度的公式很简单：在情境 c 选了做法 a 的概率是 π(a | c) = softmax(z)[a]，
        ∂ log π(a | c) / ∂ z_j = 1[j = a] − π(j | c)
（temperature = 1 时；本讲训练和评估都用 temperature = 1，即模型"本来的"分布。）

-----------------------------------------------------------------------------
【用法】
-----------------------------------------------------------------------------
    import random
    from trainable_llm import TrainableLLM, make_problem_set, evaluate

    model = TrainableLLM(1e9)
    problems = make_problem_set(random.Random(0), 100, 3, 8)
    print(evaluate(model, problems, random.Random(1), n=4))
"""

import math
import os
import random
import sys

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

from toy_llm import (ALPHA, BETA, LOGIT_BUG, LOGIT_CARELESS, PROCEDURES, REF_SIZE, ToyLLM, apply_op,  # noqa: E402
                     bar, fmt_size, make_problem, make_problem_set, mean, pad, pass_at_k, step_hardness)

TEMPERATURE = 1.0


def context_key(op, a, b):
    """
    一步运算的"情境"：(运算类型, 难度)。难度用 step_hardness 算出，保留一位小数（避免浮点误差产生不同的键）。

    【例子】47 × 12：("mul", 1.5)（乘两位数）；472 × 12：("mul", 2.2)（乘两位数 + 被乘数 ≥ 100）。
    """
    return op, round(step_hardness(op, a, b), 1)


class TrainableLLM(ToyLLM):
    """
    可以训练的 ToyLLM：在原来的 logits 上加一个按情境查表的偏移量 Δ。

    【参数】size：参数量（决定原来的 logits）；delta：{情境: [Δ_correct, Δ_bug, Δ_careless]}，默认全 0
    【Python 小知识】class TrainableLLM(ToyLLM)：继承。TrainableLLM 自动拥有 ToyLLM 的所有方法（generate、
        step_probs……），这里只**重写**（override）了 step_logits —— generate 内部调用 self.step_probs，
        self.step_probs 又调用 self.step_logits，于是自动用上了新的 logits。
    """

    def __init__(self, size, delta=None, name=None):
        super().__init__(size, cot=True, name=name)           # 【Python 小知识】super() 调用父类的 __init__
        self.delta = {k: list(v) for k, v in (delta or {}).items()}

    def step_logits(self, op, a, b, step_index):
        base = super().step_logits(op, a, b, step_index)
        d = self.delta.get(context_key(op, a, b))
        return base if d is None else [z + dz for z, dz in zip(base, d)]

    def copy(self, name=None):
        """复制一个参数相同的新模型（训练时常要保留"旧策略"或"参考策略"）。"""
        return TrainableLLM(self.size, self.delta, name or self.name)

    def p_correct(self, key):
        """
        在某个情境下选 correct 的概率（只看情境，不需要具体的数字）。用来观察训练学到了什么。

        【思路】和 ToyLLM.step_logits 的公式相同：correct 的 logit = ALPHA[运算] + BETA × log10(参数量 / 1e9) − 难度，
            bug 和 careless 的 logit 是常数，再加上这个情境的 Δ。
        """
        op, hard = key
        z = [ALPHA[op] + BETA * math.log10(self.size / REF_SIZE) - hard, LOGIT_BUG, LOGIT_CARELESS]
        d = self.delta.get(key, [0.0, 0.0, 0.0])
        z = [zi + di for zi, di in zip(z, d)]
        m = max(z)
        e = [math.exp(x - m) for x in z]
        return e[0] / math.fsum(e)


# =============================================================================
# 训练用的小工具
# =============================================================================
def step_records(problem, sample):
    """
    把一份解答拆成每一步的 (情境, 选择的做法的下标)。

    【思路】每一步的输入是上一步**写下的**值（不是正确值），所以要从 sample.values 重建每一步的输入。
    【返回】list of (key, action_index)
    """
    records = []
    x = problem.start
    for step, v, proc in zip(problem.steps, sample.values, sample.procs):
        records.append((context_key(step.op, x, step.operand), PROCEDURES.index(proc)))
        x = v
    return records


def probs_for(model, problem, sample):
    """每一步在当前模型下选那个做法的概率分布（list of 3-list），和 step_records 一一对应。"""
    out = []
    x = problem.start
    for i, (step, v) in enumerate(zip(problem.steps, sample.values)):
        out.append(model.step_probs(step.op, x, step.operand, i, TEMPERATURE))
        x = v
    return out


def add_grad(grads, key, action, probs, weight):
    """
    累加一步的梯度：weight × ∂ log π(action | key) / ∂ z = weight × (onehot − probs)。

    【参数】grads：{情境: [g0, g1, g2]}（原地修改）；weight：这一步的权重（例如 advantage）
    """
    g = grads.setdefault(key, [0.0, 0.0, 0.0])
    for j in range(3):
        g[j] += weight * ((1.0 if j == action else 0.0) - probs[j])


def apply_grads(model, grads, lr):
    """梯度上升：Δ ← Δ + lr × 梯度。"""
    for key, g in grads.items():
        d = model.delta.setdefault(key, [0.0, 0.0, 0.0])
        for j in range(3):
            d[j] += lr * g[j]


def sft(model, data, epochs=30, lr=2.0):
    """
    监督微调（SFT）：最大化 Σ log π(做法 | 情境)，数据是一批 (题目, 解答)。

    【思路】这就是"让模型更像这些解答"。每个 epoch 用全部数据算一次平均梯度，做一步梯度上升。
        STaR / rejection sampling fine-tuning 的"微调"就是它：数据只包含**做对了**的解答。
    【参数】model：被训练的模型（原地修改）；data：list of (problem, sample)；epochs、lr
    【边界】数据为空时什么也不做。
    """
    recs = [(r, problem, sample) for problem, sample in data for r in step_records(problem, sample)]
    if not recs:
        return model
    for _ in range(epochs):
        grads = {}
        for problem, sample in data:
            for (key, action), probs in zip(step_records(problem, sample), probs_for(model, problem, sample)):
                add_grad(grads, key, action, probs, 1.0 / len(recs))
        apply_grads(model, grads, lr)
    return model


def evaluate(model, problems, rng, n=4, ks=(1,)):
    """
    评估：每道题采样 n 次，返回 {k: 平均 pass@k}（无偏估计）。

    【边界】k 不能超过 n。
    """
    counts = [sum(model.generate(p, rng, TEMPERATURE).correct for _ in range(n)) for p in problems]
    return {k: mean(pass_at_k(n, c, k) for c in counts) for k in ks}


def entropy_of(model, problems, rng, n=2):
    """
    策略的平均熵（只在模型自己生成的解答所经过的情境上算）：H = −Σ π log π，单位是 nats。

    【思路】熵衡量"模型还有多不确定"：三种做法等概率时 H = log 3 ≈ 1.10；几乎确定时 H ≈ 0。
        RL 训练中熵下降得太快（entropy collapse）意味着模型不再探索 —— DAPO 要解决的问题之一。
    """
    hs = []
    for p in problems:
        for _ in range(n):
            s = model.generate(p, rng, TEMPERATURE)
            for probs in probs_for(model, p, s):
                hs.append(-math.fsum(q * math.log(q) for q in probs if q > 0))
    return mean(hs)


HARD_CONTEXTS = (("加法，无进位", ("add", 0.0)), ("乘一位数", ("mul", 0.0)),
                 ("乘两位数", ("mul", 1.5)), ("乘两位数、被乘数 ≥ 100", ("mul", 2.2)))


def context_report(model):
    """几个代表性情境下选 correct 的概率，用来看训练"教会了模型什么"。"""
    return {name: model.p_correct(key) for name, key in HARD_CONTEXTS}


if __name__ == "__main__":
    model = TrainableLLM(1e9)
    print("训练前（Δ = 0，和第 1 讲的 1B ToyLLM 相同），几个情境下选 correct 的概率：")
    for name, p in context_report(model).items():
        print(f"  {pad(name, 26)} {p:.3f}")
    problems = make_problem_set(random.Random(0), 100, 3, 8)
    print("100 道 3~8 步题上的 pass@1（temperature = 1）：", f"{evaluate(model, problems, random.Random(1))[1]:.3f}")