"""Parts 2 and 3: the KL leash, an asymmetric clipped loss, and entropy.

Three small functions; the RL loop that uses them is `rl_train` in common.py.
Test with:  python -m pytest -q -k task2   (and -k task3, -k task4)
"""

import torch
import torch.nn.functional as F


def kl_penalized_rewards(rewards, logp, ref_logp, mask, beta):
    """Task 2. Fold a KL penalty into each response's reward (the InstructGPT-style placement).

        r_i  <-  r_i - beta * sum_t mask_it * (log pi(y_it) - log pi_ref(y_it))

    The sum is the single-sample (k1) estimate of the sequence KL(pi || pi_ref)
    from Chapter 11. A response that the policy now likes much more than the reference
    did pays a price, so the optimizer only moves far from pi_ref when the reward gain
    is worth it.

    Args:
        rewards:  (B,) float tensor, one reward per response.
        logp:     (B, T) log-probs of the sampled tokens under the current policy.
        ref_logp: (B, T) log-probs of the same tokens under the frozen reference.
        mask:     (B, T) 1.0 for real response tokens, 0.0 for padding.
        beta:     float >= 0.
    Returns:
        (B,) tensor of penalized rewards.
    """
    # TODO: masked sum over t of (logp - ref_logp), then rewards - beta * that.
    raise NotImplementedError("Task 2: kl_penalized_rewards")


def clipped_pg_loss(logp, old_logp, adv, mask, eps_low=0.2, eps_high=0.2):
    """Task 3. PPO/GRPO clipped surrogate with separate lower and upper clip ranges, as a loss.

        rho_t = exp(logp_t - old_logp_t)
        obj_t = min(rho_t * A, clip(rho_t, 1 - eps_low, 1 + eps_high) * A)
        loss  = - sum_t mask_t * obj_t / sum_t mask_t        (token-level mean)

    eps_high = eps_low is the symmetric PPO clip of Chapter 9; eps_high > eps_low is
    DAPO's "clip-higher", which lets unlikely tokens with positive advantage grow more.

    Args:
        logp:     (B, T) current log-probs (with grad).
        old_logp: (B, T) log-probs under the policy that sampled the batch (no grad).
        adv:      (B,) one advantage per response, shared by all of its tokens.
        mask:     (B, T).
    Returns:
        scalar loss tensor (minimize it).
    """
    # TODO: ratio, the min of unclipped and clipped terms (broadcast adv to (B, 1)),
    #       masked token mean, minus sign.
    raise NotImplementedError("Task 3: clipped_pg_loss")


def token_entropy(logits, mask):
    """Task 4. Mean per-token entropy (nats) of the policy over response positions.

        H_t = - sum_v p_t(v) log p_t(v),   p_t = softmax(logits_t)
        return sum_t mask_t H_t / sum_t mask_t

    Use F.log_softmax for numerical stability (p log p with p = 0 must give 0, not nan).

    Args:
        logits: (B, T, V).
        mask:   (B, T).
    Returns:
        scalar tensor.
    """
    # TODO
    raise NotImplementedError("Task 4: token_entropy")