"""Tested loss primitives, not a claim of completed large-scale RL training."""
import torch
from torch.nn import functional as F


def masked_sft_loss(logits, labels, assistant_mask):
    if labels.shape != assistant_mask.shape or logits.shape[:-1] != labels.shape:
        raise ValueError("Incompatible SFT shapes")
    if not assistant_mask.any():
        raise ValueError("No assistant target tokens")
    targets = labels.masked_fill(~assistant_mask.bool(), -100)
    return F.cross_entropy(logits.flatten(0, -2), targets.flatten(), ignore_index=-100)


def dpo_loss(policy_chosen, policy_rejected, reference_chosen, reference_rejected, beta=0.1):
    if beta <= 0:
        raise ValueError("beta must be positive")
    shapes = {x.shape for x in (policy_chosen, policy_rejected, reference_chosen, reference_rejected)}
    if len(shapes) != 1:
        raise ValueError("Log-probability shapes must match")
    margin = (policy_chosen-policy_rejected) - (reference_chosen-reference_rejected)
    return -F.logsigmoid(beta * margin).mean()


def group_advantages(rewards, eps=1e-6):
    if rewards.ndim != 2 or rewards.shape[1] < 2 or not torch.isfinite(rewards).all():
        raise ValueError("Expected finite batch x group rewards, group size >= 2")
    return (rewards-rewards.mean(-1, keepdim=True)) / rewards.std(-1, keepdim=True, unbiased=False).clamp_min(eps)


def rejection_sample(candidates, verifier):
    """Return only independently verified candidates; never reward persuasive prose."""
    accepted = []
    for candidate in candidates:
        try:
            if verifier(candidate) is True:
                accepted.append(candidate)
        except Exception:
            continue
    return accepted