"""Client-side RL losses for forward_backward_custom (sum over loss tokens, Tinker convention).

grpo: Spindle stock ppo (modal-projects/spindle@e58ebef, megatron_runtime/common/forward_backward.py::_loss):
      loss = -sum(mask * min(r*A, clamp(r, 1-eps, 1+eps)*A)), r = exp(logp - sampler_logp)
dppo: Stable-RL DPPO-Binary-TV (sail-sg/Stable-RL@c6af1d1, verl/trainer/ppo/core_algos.py, loss_mode=dppo_binary_tv):
      invalid = A>0 & (p - mu) > delta  |  A<0 & (p - mu) < -delta   (mu = sampler prob)
      loss = -sum(mask * A * sg(min(r, 10)) * keep * logp), r = exp(min(logp - sampler_logp, 20))
Both arms compute every diagnostic on the same tensors so masks/clips are observable in either arm.
"""
import torch

GRPO_EPS = 0.2
DPPO_DELTA = 0.15
DPPO_CLIP_RATIO_C = 10.0


def token_terms(logp, old_lp, adv, mask, mode, eps=GRPO_EPS, delta=DPPO_DELTA):
    if mode == "grpo":
        ratio = torch.exp(logp - old_lp)
        obj = torch.minimum(ratio * adv, torch.clamp(ratio, 1 - eps, 1 + eps) * adv)
        return -(obj * mask)
    if mode == "dppo":
        prob, rollout_prob = torch.exp(logp), torch.exp(old_lp)
        invalid_pos = (prob - rollout_prob) > delta
        invalid_neg = (prob - rollout_prob) < -delta
        invalid = torch.where(adv > 0, invalid_pos, invalid_neg)
        keep = (1.0 - invalid.detach().float()) * mask
        ratio = torch.clamp(torch.exp(torch.clamp(logp - old_lp, max=20.0)), max=DPPO_CLIP_RATIO_C).detach()
        return -adv * ratio * keep * logp
    raise ValueError(mode)


@torch.no_grad()
def diagnostics(logp, old_lp, adv, mask, eps=GRPO_EPS, delta=DPPO_DELTA):
    m = mask > 0
    n = int(m.sum())
    if n == 0:
        return {}
    lp, olp, a = logp[m].double(), old_lp[m].double(), adv[m].double()
    p, mu, r = lp.exp(), olp.exp(), (lp - olp).exp()
    dppo_masked = torch.where(a > 0, (p - mu) > delta, (p - mu) < -delta)
    grpo_clipped = ((a > 0) & (r > 1 + eps)) | ((a < 0) & (r < 1 - eps))
    absdp = (p - mu).abs()
    return {
        "n_loss_tokens": float(n),
        "dppo_masked_token_fraction": float(dppo_masked.double().mean()),
        "grpo_clip_fraction": float(grpo_clipped.double().mean()),
        "policy_drift_abs_p_q90": float(torch.quantile(absdp.float(), 0.9)),
        "policy_drift_abs_p_mean": float(absdp.mean()),
        "mean_ratio": float(r.mean()),
        "nonzero_adv_fraction": float((a != 0).double().mean()),
    }


def make_loss(mode, sink=None):
    """Returns CustomLossFnV1: (data, logprobs_list) -> (loss, metrics). Datums carry
    loss_fn_inputs logprobs (sampler), advantages, mask (action tokens)."""

    def loss_fn(data, logprobs_list):
        total = 0.0
        cat = {k: [] for k in ("lp", "olp", "adv", "mask")}
        for datum, logp in zip(data, logprobs_list, strict=True):
            inp = datum.loss_fn_inputs
            old_lp = inp["logprobs"].to_torch().to(logp.dtype)
            adv = inp["advantages"].to_torch().to(logp.dtype)
            mask = inp["mask"].to_torch().to(logp.dtype)
            total = total + token_terms(logp, old_lp, adv, mask, mode).sum()
            for k, v in zip(cat, (logp.detach(), old_lp, adv, mask)):
                cat[k].append(v)
        diag = diagnostics(*(torch.cat(cat[k]) for k in cat))
        if sink is not None:
            sink.append(diag)
        return total, {f"{k}": v for k, v in diag.items()}

    return loss_fn