"""第 28 章：两步算式环境中的策略梯度与状态价值，不训练语言模型。"""

import torch
from torch import nn


def discounted_returns(rewards, gamma):
    result, future = [], 0.0
    for reward in reversed(rewards):
        future = reward+gamma*future
        result.append(future)
    return list(reversed(result))


def expected_reward(logits):
    probabilities = logits.softmax(-1)
    # 只有先 2+3、再乘 4 的路径得到 1 分。
    return probabilities[0, 0]*probabilities[1, 0]


def train(steps=200):
    torch.set_num_threads(1)
    generator = torch.Generator().manual_seed(2026)
    logits = nn.Parameter(torch.zeros(3, 2, dtype=torch.float64))
    values = nn.Parameter(torch.zeros(3, dtype=torch.float64))
    optimizer = torch.optim.SGD([{"params": [logits], "lr": .2},
                                {"params": [values], "lr": .1}])
    history = [(0, expected_reward(logits).item())]
    for step in range(1, steps+1):
        probabilities = logits.softmax(-1)
        first = torch.multinomial(probabilities[0].detach(), 32, replacement=True, generator=generator)
        second_states = 1+first  # 状态 1 的中间结果为 5，状态 2 的中间结果为 6。
        second = torch.multinomial(probabilities[second_states].detach(), 1, generator=generator).squeeze(1)
        states = torch.stack([torch.zeros_like(first), second_states], dim=1)
        actions = torch.stack([first, second], dim=1)
        final_results = torch.where(second == 0, (5+first)*4, (5+first)+4)
        rewards = (final_results == 20).double()
        # gamma=1；两步都使用同一个终局回报。价值作为基线时停止梯度。
        returns = rewards[:, None].expand(-1, 2)
        advantages = returns-values.detach()[states]
        chosen_log_probs = logits.log_softmax(-1)[states, actions]
        actor_loss = -(advantages*chosen_log_probs).sum(1).mean()
        value_loss = .5*((values[states]-returns)**2).mean()
        optimizer.zero_grad(set_to_none=True)
        (actor_loss+value_loss).backward()
        optimizer.step()
        if step % 20 == 0:
            history.append((step, expected_reward(logits).item()))
    return logits.detach(), values.detach(), history


def main():
    logits, values, history = train()
    probabilities = logits.softmax(-1)
    print("任务：先合并 2 和 3，再与 4 运算；最终结果 20 得 1 分。")
    print(f"训练前成功概率={history[0][1]:.3f}；训练后={history[-1][1]:.3f}")
    print(f"第一步选择加法的概率={probabilities[0, 0]:.3f}")
    print(f"得到 5 后选择乘法的概率={probabilities[1, 0]:.3f}")
    print("学得的状态价值：" + ", ".join(f"{value:.3f}" for value in values))
    exact = [expected_reward(logits), probabilities[1, 0], torch.tensor(0.0)]
    print("当前策略的精确状态价值：" + ", ".join(f"{value:.3f}" for value in exact))
    print("奖励 [0, 0, 1]、gamma=0.9 的回报：" +
          ", ".join(f"{value:.2f}" for value in discounted_returns([0, 0, 1], .9)))
    print("每 20 次更新的成功概率：" + ", ".join(f"{value:.3f}" for _, value in history))


if __name__ == "__main__":
    main()