"""第 30 章：有限候选上的结果验证、过程验证与组相对更新；不训练语言模型。"""

import re
import torch
from torch import nn
from torch.nn import functional as F
from preference_optimization import clipped_gain

TRACES = [
    "2+3=5;5*4=20;答案=20",
    "2+3=5;5+4=9;答案=9",
    "2+3=6;6*4=20;答案=20",
    "答案=20或9",
]


def outcome_reward(text):
    # 只读取严格格式的最后一段，不评判之前的推导。
    final = text.split(";")[-1]
    match = re.fullmatch(r"答案=(-?\d+)", final)
    return int(match is not None and int(match.group(1)) == 20)


def process_reward(text):
    parts = text.split(";")
    if len(parts) != 3 or not outcome_reward(text):
        return 0
    parsed = [re.fullmatch(r"(-?\d+)([+*])(-?\d+)=(-?\d+)", part) for part in parts[:2]]
    if not all(parsed):
        return 0
    first, second = [match.groups() for match in parsed]
    a, op1, b, middle = int(first[0]), first[1], int(first[2]), int(first[3])
    left, op2, c, result = int(second[0]), second[1], int(second[2]), int(second[3])
    return int((a, b, c, op1, op2) == (2, 3, 4, "+", "*")
               and middle == a+b and left == middle and result == left*c)


def advantages(rewards):
    # 使用总体标准差；同分组的奖励项梯度为零。
    return (rewards-rewards.mean())/(rewards.std(unbiased=False)+1e-8)


def train(verifier, rounds=160, group_size=8, beta=.05):
    logits = nn.Parameter(torch.zeros(4, dtype=torch.float64))
    optimizer = torch.optim.SGD([logits], lr=.15)
    generator = torch.Generator().manual_seed(2026)
    reference = torch.full((4,), .25, dtype=torch.float64)
    all_rewards = torch.tensor([verifier(trace) for trace in TRACES], dtype=torch.float64)
    history, equal_groups = [], 0
    for step in range(1, rounds+1):
        old = logits.detach().softmax(-1)
        chosen = torch.multinomial(old, group_size, replacement=True, generator=generator)
        sampled_rewards = all_rewards[chosen]
        equal_groups += bool((sampled_rewards == sampled_rewards[0]).all())
        relative = advantages(sampled_rewards).detach()
        old_log = old[chosen].log()
        for _ in range(3):
            log_probs = logits.log_softmax(-1)
            ratio = (log_probs[chosen]-old_log).exp()
            surrogate = clipped_gain(ratio, relative).mean()
            kl = (log_probs.exp()*(log_probs-reference.log())).sum()
            loss = -surrogate+beta*kl
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
        if step % 20 == 0:
            probabilities = logits.detach().softmax(-1)
            history.append((step, float(probabilities @ all_rewards), float(probabilities[0])))
    return logits.detach().softmax(-1), history, equal_groups


def distill_filtered(teacher):
    generator = torch.Generator().manual_seed(2027)
    candidates = torch.multinomial(teacher, 32, replacement=True, generator=generator)
    outcome_ids = [i for i in candidates.tolist() if outcome_reward(TRACES[i])]
    process_ids = [i for i in candidates.tolist() if process_reward(TRACES[i])]
    if not process_ids:
        raise RuntimeError("没有通过过程检查的样本，不能开始本例蒸馏")
    student = nn.Parameter(torch.zeros(4, dtype=torch.float64))
    optimizer = torch.optim.SGD([student], lr=.2)
    targets = torch.tensor(process_ids)
    for _ in range(100):
        loss = F.cross_entropy(student.expand(len(targets), -1), targets)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()
    return len(outcome_ids), len(process_ids), student.detach().softmax(-1)


def main():
    torch.set_num_threads(1)
    for i, trace in enumerate(TRACES):
        print(f"候选 {i}：结果奖励={outcome_reward(trace)}，过程奖励={process_reward(trace)}；{trace}")
    outcome, _, equal_outcome = train(outcome_reward)
    process, _, equal_process = train(process_reward)
    print("结果奖励训练后的候选概率：" + ", ".join(f"{p:.3f}" for p in outcome))
    print("过程奖励训练后的候选概率：" + ", ".join(f"{p:.3f}" for p in process))
    print(f"160 轮中的同分组：结果={equal_outcome}；过程={equal_process}")
    print("同分组 [1,1,1,1] 的优势：" +
          ", ".join(f"{a:.1f}" for a in advantages(torch.ones(4, dtype=torch.float64))))
    count_outcome, count_process, student = distill_filtered(outcome)
    print(f"从结果奖励策略采样 32 份：结果通过={count_outcome}，过程通过={count_process}")
    print(f"只用过程通过样本训练学生后，正确候选概率={student[0]:.3f}")
    for n in [1, 4, 8]:
        print(f"均匀采样 {n} 次：至少一份结果通过={1-.5**n:.3f}；"
              f"至少一份过程通过={1-.75**n:.3f}")


if __name__ == "__main__":
    main()