"""第 29 章：三候选回答上的奖励拟合、DPO 与单步 PPO；不是完整 LLM RLHF。"""

import torch
from torch import nn
from torch.nn import functional as F

REFERENCE = torch.tensor([.2, .6, .2], dtype=torch.float64)
RESPONSES = ["红", "红色。", "蓝"]


def preference_pairs():
    # 同一个提示：两份正确答案比较，8 个模拟标注选简短，2 个选完整句。
    winners = torch.tensor([0]*8+[1]*2)
    losers = 1-winners
    return winners, losers


def dpo_loss(logits, reference, winners, losers, beta=1):
    log_probs, log_ref = logits.log_softmax(-1), reference.log()
    margin = beta*((log_probs[winners]-log_ref[winners])
                   -(log_probs[losers]-log_ref[losers]))
    return -F.logsigmoid(margin).mean()


def clipped_gain(ratio, advantage, epsilon=.2):
    return torch.minimum(ratio*advantage, ratio.clamp(1-epsilon, 1+epsilon)*advantage)


def fit_reward_and_dpo():
    winners, losers = preference_pairs()
    rewards = nn.Parameter(torch.zeros(3, dtype=torch.float64))
    logits = nn.Parameter(REFERENCE.log().clone())
    optimizer = torch.optim.SGD([rewards, logits], lr=.3)
    for _ in range(300):
        reward_loss = -F.logsigmoid(rewards[winners]-rewards[losers]).mean()
        loss = reward_loss+dpo_loss(logits, REFERENCE, winners, losers)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()
    return rewards.detach(), logits.detach()


def ppo_bandit(rewards, rounds=100):
    logits = nn.Parameter(REFERENCE.log().clone())
    optimizer = torch.optim.SGD([logits], lr=.2)
    generator = torch.Generator().manual_seed(2026)
    for _ in range(rounds):
        # old 随每轮采样更新；REFERENCE 在整个实验中固定。
        old_probs = logits.detach().softmax(-1)
        actions = torch.multinomial(old_probs, 64, replacement=True, generator=generator)
        old_log_probs = old_probs[actions].log()
        advantages = (rewards[actions]-(old_probs*rewards).sum()).detach()
        for _ in range(4):
            log_probs = logits.log_softmax(-1)
            ratio = (log_probs[actions]-old_log_probs).exp()
            gain = clipped_gain(ratio, advantages).mean()
            # 三个候选全部可枚举，所以用精确 KL，不使用采样近似。
            kl = (log_probs.exp()*(log_probs-REFERENCE.log())).sum()
            loss = -gain+kl
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
    return logits.detach()


def overoptimization():
    # 独立假设：第三个错误答案被一个有缺陷的代理奖励高估。
    proxy = torch.tensor([1., .8, 3.], dtype=torch.float64)
    utility = torch.tensor([1., 1., 0.], dtype=torch.float64)
    rows = []
    for beta in [2., 1., .2]:
        probabilities = (REFERENCE.log()+proxy/beta).softmax(-1)
        rows.append((beta, float(probabilities @ proxy), float(probabilities @ utility)))
    return rows


def main():
    torch.set_num_threads(1)
    rewards, dpo = fit_reward_and_dpo()
    ppo = ppo_bandit(rewards)
    print("同一提示的三个候选：" + " | ".join(RESPONSES))
    print("模拟偏好：简短正确答案赢 8 次，完整正确答案赢 2 次；错误答案未参与比较。")
    for name, probabilities in [("参考策略", REFERENCE), ("DPO", dpo.softmax(-1)),
                                ("单步 PPO", ppo.softmax(-1)),
                                ("拟合奖励的解析最优", (REFERENCE.log()+rewards).softmax(-1))]:
        print(name + "：" + ", ".join(f"{value:.3f}" for value in probabilities))
    print(f"奖励差={rewards[0]-rewards[1]:.3f}；"
          f"拟合的简短偏好概率={torch.sigmoid(rewards[0]-rewards[1]):.3f}")
    ratios = torch.tensor([.7, 1., 1.3], dtype=torch.float64)
    print("ratio=0.7,1.0,1.3；正优势的裁剪收益：" +
          ", ".join(f"{value:.1f}" for value in clipped_gain(ratios, torch.ones(3))))
    print("ratio=0.7,1.0,1.3；负优势的裁剪收益：" +
          ", ".join(f"{value:.1f}" for value in clipped_gain(ratios, -torch.ones(3))))
    print("另一个假设：代理奖励高估错误答案。")
    for beta, proxy, utility in overoptimization():
        print(f"beta={beta:.1f}：代理奖励={proxy:.3f}；真实正确概率={utility:.3f}")


if __name__ == "__main__":
    main()