"""Chapter 21 lab, Tasks 1-3: on-policy distillation (OPD).

The student pi_theta samples a response y to prompt x. A frozen teacher pi_T then
scores every token the *student* produced. The student is pushed toward the teacher
on exactly the states it actually visits (the DAgger idea from Chapter 10).

Two ways to turn "be like the teacher" into a gradient, both used in 2026 reports:

  * sampled-token OPD (Tasks 1-2; Thinking Machines blog, Kimi K3): treat
        r_t = log pi_T(y_t | s_t) - log pi_theta(y_t | s_t)
    as a per-token reward and plug it into an ordinary policy-gradient loss with
    discount 0 (each token is rewarded only for itself).
  * full-vocabulary OPD (Task 3; DeepSeek-V4): at every visited state s_t compute
    the exact KL( pi_theta(.|s_t) || pi_T(.|s_t) ) over the whole vocabulary and
    minimize it directly.

Shapes: B = batch of responses, T = response positions (tinylm's MAX_ANSWER = 4),
V = vocabulary size. `mask[b, t] = 1` for tokens the student actually generated
(including <eos>), 0 for padding after <eos>.

Tests: test_opd.py. The full experiment: run.py.
"""

import torch
import torch.nn.functional as F  # noqa: F401  (you will want log_softmax)


# ---------------------------------------------------------------- Task 1
def reverse_kl_rewards(student_logp, teacher_logp, mask):
    """Per-token OPD reward: minus the single-sample estimate of the reverse KL.

        r_t = log pi_T(y_t | s_t) - log pi_theta(y_t | s_t)        (masked tokens -> 0)

    E_{y_t ~ pi_theta}[-r_t] = KL(pi_theta(.|s_t) || pi_T(.|s_t)), so a policy that
    maximizes r_t is pulled toward the teacher. r_t > 0 means "the teacher likes this
    token more than you do: do it more".

    Parameters
    ----------
    student_logp : (B, T) float tensor, log pi_theta of the sampled tokens (may carry grad).
    teacher_logp : (B, T) float tensor, log pi_T of the same tokens.
    mask         : (B, T) float tensor of 0/1.

    Returns
    -------
    (B, T) float tensor with NO gradient attached (it is a reward, not a loss),
    equal to (teacher_logp - student_logp) on masked-in tokens and 0 elsewhere.
    """
    # TODO: difference of log-probs, times the mask, detached from the graph.
    raise NotImplementedError("Task 1: implement reverse_kl_rewards")


# ---------------------------------------------------------------- Task 2
def opd_pg_loss(student_logp, rewards, mask):
    """REINFORCE-style loss for sampled-token OPD with discount 0.

    The advantage of token t is just its own reward, A_t = r_t (no reward-to-go,
    no baseline, no importance ratio: we take one gradient step per fresh batch,
    so the data is exactly on-policy). Average over all generated tokens in the
    batch ("token-mean"):

        loss = - sum_{b,t} mask[b,t] * sg(r[b,t]) * student_logp[b,t]  /  sum_{b,t} mask[b,t]

    where sg = stop-gradient (treat r as a constant).

    Parameters
    ----------
    student_logp : (B, T) tensor WITH grad, log pi_theta(y_t | s_t).
    rewards      : (B, T) tensor, output of reverse_kl_rewards (detach it anyway).
    mask         : (B, T) float 0/1 tensor.

    Returns
    -------
    Scalar tensor. Its gradient w.r.t. student_logp[b, t] is -r[b,t] * mask[b,t] / sum(mask).
    """
    # TODO: one line, plus a guard so an all-zero mask does not divide by zero.
    raise NotImplementedError("Task 2: implement opd_pg_loss")


# ---------------------------------------------------------------- Task 3
def full_vocab_reverse_kl(student_logits, teacher_logits, mask):
    """Exact reverse KL at every visited position, averaged over generated tokens.

        KL_t = sum_v pi_theta(v | s_t) * ( log pi_theta(v | s_t) - log pi_T(v | s_t) )
        loss = sum_{b,t} mask[b,t] * KL_t[b,t]  /  sum_{b,t} mask[b,t]

    The states s_t were visited by the *student* (that is what makes it on-policy);
    we do not differentiate through which states got visited, only through the
    student's distribution at each of them.

    Parameters
    ----------
    student_logits : (B, T, V) tensor WITH grad.
    teacher_logits : (B, T, V) tensor (treat as constant: detach it).
    mask           : (B, T) float 0/1 tensor.

    Returns
    -------
    Scalar tensor >= 0, zero when the two distributions agree at every masked position.
    Use log_softmax (not log(softmax)) for numerical stability.
    """
    # TODO: log_softmax both; KL per position; masked mean.
    raise NotImplementedError("Task 3: implement full_vocab_reverse_kl")