"""Original toy GRPO update and editorial R1 process lesson.

One prompt feature vector, four synthetic one-step completion IDs, fixed old and
reference policies. This is loss arithmetic, not a language/reasoning model.
Population std with epsilon handles equal-reward groups; clip epsilon .2, KL .04.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F


class CompletionPolicy(nn.Module):
    def __init__(self):
        super().__init__()
        self.logits = nn.Linear(3, 4)
        self.log_probabilities = nn.LogSoftmax(-1)

    def forward(self, prompt, completions):
        return self.log_probabilities(self.logits(prompt))[0, completions]


class GroupAdvantages(nn.Module):
    def forward(self, rewards):
        return (rewards - rewards.mean()) / (rewards.std(correction=0) + 1e-8)


class PolicyRatio(nn.Module):
    def forward(self, current, old):
        return (current - old).exp()


class ClippedSurrogate(nn.Module):
    def forward(self, ratio, advantages):
        return torch.minimum(ratio * advantages, ratio.clamp(0.8, 1.2) * advantages)


class ReferencePenalty(nn.Module):
    def forward(self, current, reference):
        difference = reference - current
        return difference.exp() - difference - 1


class ToyGRPO(nn.Module):
    def __init__(self):
        super().__init__()
        self.policy = CompletionPolicy()
        self.old_policy = CompletionPolicy()
        self.reference_policy = CompletionPolicy()
        self.old_policy.load_state_dict(self.policy.state_dict())
        self.reference_policy.load_state_dict(self.policy.state_dict())
        for parameter in self.old_policy.parameters():
            parameter.requires_grad_(False)
        for parameter in self.reference_policy.parameters():
            parameter.requires_grad_(False)
        self.advantages = GroupAdvantages()
        self.ratio = PolicyRatio()
        self.clipped = ClippedSurrogate()
        self.divergence = ReferencePenalty()

    def forward(self, prompt, completions, rewards):
        current = self.policy(prompt, completions)
        old = self.old_policy(prompt, completions)
        reference = self.reference_policy(prompt, completions)
        advantages = self.advantages(rewards)
        ratio = self.ratio(current, old)
        surrogate = self.clipped(ratio, advantages)
        penalty = self.divergence(current, reference)
        return -(surrogate - 0.04 * penalty).mean().reshape(1)


def setup():
    torch.manual_seed(71)
    model = ToyGRPO().double()
    prompt = torch.tensor([[0.5, -0.3, 1.0]], dtype=torch.float64)
    completions = torch.arange(4)
    rewards = torch.tensor([0.0, 1.0, 0.0, 0.5], dtype=torch.float64)
    return model, prompt, completions, rewards


def verify():
    model, prompt, ids, rewards = setup()
    advantages = model.advantages(rewards)
    torch.testing.assert_close(
        advantages.mean(), torch.tensor(0.0, dtype=torch.float64)
    )
    torch.testing.assert_close(model.advantages(rewards + 7), advantages)
    torch.testing.assert_close(
        model.advantages(rewards * 2), advantages, atol=2e-8, rtol=2e-8
    )
    assert torch.count_nonzero(model.advantages(torch.ones(4))) == 0
    ratios = torch.tensor([1.5, 0.5, 1.5, 0.5], dtype=torch.float64, requires_grad=True)
    adv = torch.tensor([1.0, -1.0, -1.0, 1.0], dtype=torch.float64)
    surrogate = model.clipped(ratios, adv)
    torch.testing.assert_close(
        surrogate, torch.tensor([1.2, -0.8, -1.5, 0.5], dtype=torch.float64)
    )
    torch.testing.assert_close(
        torch.autograd.grad(surrogate.sum(), ratios)[0],
        torch.tensor([0.0, 0.0, -1.0, 1.0], dtype=torch.float64),
    )
    current = model.policy(prompt, ids)
    reference = model.reference_policy(prompt, ids)
    torch.testing.assert_close(
        model.divergence(current, reference), torch.zeros(4, dtype=torch.float64)
    )
    assert torch.all(
        model.divergence(current + torch.tensor([1.0, -1.0, 2.0, -2.0]), reference) >= 0
    )
    opt = torch.optim.SGD(model.policy.parameters(), lr=0.05)
    before = current[1].exp().item()
    old = model.old_policy(prompt, ids).clone()
    opt.zero_grad()
    model(prompt, ids, rewards).sum().backward()
    assert all(p.grad is None for p in model.old_policy.parameters())
    assert all(p.grad is None for p in model.reference_policy.parameters())
    opt.step()
    assert model.policy(prompt, ids)[1].exp().item() > before
    torch.testing.assert_close(model.old_policy(prompt, ids), old, atol=0, rtol=0)
    return [
        "Group advantages center rewards, preserve reward shifts/scales and vanish for equal rewards",
        "Clipping handles positive and negative advantages with the expected saturated and unsaturated gradients",
        "The sampled reference penalty is nonnegative and zero for identical policies",
        "One toy GRPO update increases the best-rewarded completion probability while old/reference policies remain frozen",
    ]


R1_STAGES = [
    dict(
        label="Cold-start SFT",
        detail="Fine-tune DeepSeek-V3-Base on thousands of readable long-response examples.",
    ),
    dict(
        label="Reasoning-oriented RL",
        detail="Apply GRPO with reasoning rewards and language consistency; the policy generates data for later curation.",
    ),
    dict(
        label="Curated SFT",
        detail="Filter about 600k reasoning samples, mix about 200k other samples, and fine-tune DeepSeek-V3-Base for two epochs.",
    ),
    dict(
        label="RL for all scenarios",
        detail="Combine rule-based reasoning rewards with preference reward models for general tasks.",
    ),
]


def experiment():
    model, prompt, ids, rewards = setup()
    cases = []
    specs = [
        (
            "zero",
            "R1-Zero · direct RL",
            "R1-Zero",
            [
                dict(
                    label="DeepSeek-V3-Base",
                    detail="Begin with a pretrained base model; no cold-start SFT.",
                ),
                dict(
                    label="Direct GRPO",
                    detail="Use accuracy and format rewards on sampled responses. This is the R1-Zero path, separate from R1.",
                ),
            ],
            1,
        ),
        ("cold", "R1 · cold-start SFT", "DeepSeek-R1", R1_STAGES, 0),
        ("rl", "R1 · reasoning-oriented RL", "DeepSeek-R1", R1_STAGES, 1),
        ("curate", "R1 · curation and SFT", "DeepSeek-R1", R1_STAGES, 2),
        ("align", "R1 · final RL", "DeepSeek-R1", R1_STAGES, 3),
        (
            "distill",
            "Distilled models · supervised transfer",
            "R1 distillation",
            [
                dict(
                    label="Smaller Qwen/Llama base",
                    detail="Start from a separate smaller dense model.",
                ),
                dict(
                    label="SFT on curated responses",
                    detail="Use the approximately 800k curated examples. The reported distilled models use SFT without a further RL stage.",
                ),
            ],
            1,
        ),
    ]
    for identity, label, process, steps, active in specs:
        current = model.policy(prompt, ids)
        old = model.old_policy(prompt, ids)
        reference = model.reference_policy(prompt, ids)
        advantages = model.advantages(rewards)
        ratios = model.ratio(current, old)
        penalty = model.divergence(current, reference)
        surrogate = model.clipped(ratios, advantages)
        with torch.no_grad():
            cases.append(
                dict(
                    id=identity,
                    label=label,
                    target="policy"
                    if identity in ["cold", "curate", "distill"]
                    else "advantages",
                    note="The highlighted stage describes the original January 2025 paper. It is an editorial process map, not execution status. Below it, the same independent four-candidate toy GRPO batch explains reward advantages, policy ratios and the reference penalty. It is not a simulation of the selected SFT/RL stage or evidence of learned reasoning.",
                    process=dict(label=process, steps=steps, active=active),
                    vectors=[
                        dict(
                            label="Synthetic completion rewards · IDs 0–3",
                            values=rewards.tolist(),
                        ),
                        dict(
                            label="Group-relative advantages",
                            values=advantages.tolist(),
                        ),
                        dict(
                            label="Current / old probability ratio",
                            values=ratios.tolist(),
                        ),
                        dict(
                            label="Sampled reference penalty", values=penalty.tolist()
                        ),
                    ],
                    metrics=[
                        dict(
                            label="Toy GRPO loss · before update",
                            value=(-(surrogate - 0.04 * penalty).mean()).item(),
                        ),
                        dict(label="Clipping epsilon", value=0.2),
                        dict(label="KL coefficient", value=0.04),
                    ],
                )
            )
    # Two additional arithmetic cases demonstrate the effect of the authorized local update.
    opt = torch.optim.SGD(model.policy.parameters(), lr=0.05)
    before = model.policy(prompt, ids).detach().exp()
    opt.zero_grad()
    model(prompt, ids, rewards).sum().backward()
    opt.step()
    with torch.no_grad():
        after = model.policy(prompt, ids).exp()
        ratios = model.ratio(model.policy(prompt, ids), model.old_policy(prompt, ids))
        cases.append(
            dict(
                id="update",
                label="Toy arithmetic · one GRPO update",
                target="clipped",
                note="One SGD step changes a four-choice toy policy. The fixed candidate group is enumerated to make the arithmetic inspectable; it is not sampled language, R1 training, or a reasoning benchmark. Old and reference policies retain their original weights.",
                process=dict(
                    label="Local arithmetic demonstration",
                    steps=[
                        dict(
                            label="Fixed toy group",
                            detail="Four synthetic completion IDs with supplied rewards.",
                        ),
                        dict(
                            label="One GRPO update",
                            detail="Only the current policy changes; old and reference snapshots stay fixed.",
                        ),
                    ],
                    active=1,
                ),
                vectors=[
                    dict(label="Probability before update", values=before.tolist()),
                    dict(label="Probability after update", values=after.tolist()),
                    dict(label="Updated current / old ratio", values=ratios.tolist()),
                ],
                metrics=[
                    dict(
                        label="Probability gain for best-rewarded ID 1",
                        value=(after[1] - before[1]).item(),
                    )
                ],
            )
        )
    return dict(
        kind="process",
        title="The milestone is a learning process.",
        description="Distinguish R1-Zero, R1’s four stages, and supervised distillation. A separate original toy GRPO implementation provides reproducible arithmetic without claiming to reproduce the published model.",
        controlLabel="Training process",
        cases=cases,
    )