"""Chapter 16 lab: one configurable loss for the whole critic-free zoo.

The chapter's master formula says every critic-free algorithm (REINFORCE, RLOO,
GRPO, Dr. GRPO, DAPO, GSPO, CISPO, REINFORCE++, ...) is the same estimator

    grad J  ~  sum_i sum_t  w_{i,t} * A_i * grad log pi_theta(y_{i,t} | x, y_{i,<t})

with different settings of six knobs. You will write that estimator once, as a
loss whose gradient is the formula above, and every preset in common.PRESETS
will then be one line of configuration.

Your tasks (in order; run `python -m pytest -q -k task1` etc. after each):

    Task 1  group_advantages   knobs 1-2: the baseline, and the scale (std or not)
    Task 2  aggregate          knob 3: seq-mean vs token-mean vs constant normaliser
    Task 3  token_surrogate    knob 4: token vs sequence ratio; PPO clip vs CISPO
    Task 4  kl_per_token       knob 5: the k1 / k2 / k3 KL estimators
    Task 5  pg_loss            all knobs together

Shapes: N responses in the batch (number of prompts x G), T = max response length.
    logp, old_logp, ref_logp, mask : (N, T) float; mask is 1.0 on real response tokens
    rewards                        : (N,) float
    group_ids                      : (N,) long; responses to the same prompt share an id 0..P-1

Notation (as in the chapter): rho = pi_theta / pi_old (importance ratio),
A = advantage, |y_i| = mask[i].sum() = length of response i.
"""

import torch

from common import PGConfig  # noqa: F401  (the knob settings; read its docstring in common.py)


# ---------------------------------------------------------------------------
# Task 1: advantages (knob 1: baseline, knob 2: scale)
# ---------------------------------------------------------------------------
def group_advantages(rewards, group_ids, baseline="group_mean", scale="group_std", eps=1e-6):
    """One advantage per response. Returns an (N,) float tensor.

    Step 1, subtract a baseline b_i (the chapter's "knob 1"):
        "none"        b_i = 0
        "batch_mean"  b_i = mean of all N rewards
        "group_mean"  b_i = mean of the rewards in i's group (including r_i itself)  (GRPO)
        "loo"         b_i = mean of the OTHER rewards in i's group:
                            (sum_group - r_i) / (n_group - 1)                    (RLOO)
                      (if a group has a single response, use b_i = 0)
        A_i = r_i - b_i

    Step 2, scale (the chapter's "knob 2"):
        "none"        A_i unchanged
        "group_std"   A_i / (std of the REWARDS in i's group + eps)              (GRPO)
        "batch_std"   (A_i - mean_j A_j) / (std_j A_j + eps), over the whole batch  (REINFORCE++)
    Use the population std (divide by n, i.e. torch.std(..., unbiased=False)).

    Worked example (two prompts, G = 4), rewards [1,0,0,1, 0,0,0,1]:
        group_mean, none : [ .5,-.5,-.5, .5,  -.25,-.25,-.25,.75]
        loo, none        : [2/3,-2/3,-2/3,2/3, -1/3,-1/3,-1/3, 1 ]
        group_mean, group_std: [1,-1,-1,1, -0.577,-0.577,-0.577,1.732]

    Hints: a loop over torch.unique(group_ids) is perfectly fine. (The reference
    solution uses index_add_, but clarity beats speed here.)
    """
    # TODO: compute the baseline, subtract it, then apply the scale.
    raise NotImplementedError("Task 1: group_advantages")


# ---------------------------------------------------------------------------
# Task 2: aggregation (knob 3)
# ---------------------------------------------------------------------------
def aggregate(per_token, mask, aggregation="seq_mean", max_len=1):
    """Turn (N, T) per-token terms into one scalar. Only masked-in tokens count.

        "seq_mean"    (1/N) sum_i [ (1/|y_i|) sum_t x_{i,t} ]        (GRPO, GSPO)
        "token_mean"  (1 / sum_i |y_i|) sum_i sum_t x_{i,t}          (DAPO, CISPO)
        "constant"    (1 / (N * max_len)) sum_i sum_t x_{i,t}        (Dr. GRPO; with
                       max_len = 1 it is (1/N) sum_i sum_t, plain REINFORCE's sequence sum)

    where |y_i| = mask[i].sum(). Guard against |y_i| = 0 with .clamp(min=1).

    Why it matters: under "seq_mean" every token of a long response gets a
    SMALLER weight than every token of a short one (the length bias of Section
    "Knob 3" in the chapter). Under the other two, every token gets the same weight.
    """
    # TODO: three cases.
    raise NotImplementedError("Task 2: aggregate")


# ---------------------------------------------------------------------------
# Task 3: importance ratio and clipping (knob 4)
# ---------------------------------------------------------------------------
def token_surrogate(logp, old_logp, adv, mask, cfg):
    """Per-token surrogate terms (N, T) whose GRADIENT is the chosen policy gradient.

    Inputs: logp (N,T) with grad; old_logp (N,T) no grad; adv (N,) (detach it!);
    mask (N,T); cfg a PGConfig (uses cfg.ratio, cfg.clip, cfg.eps_low, cfg.eps_high).

    1. The ratio. Let lr_{i,t} = (logp - old_logp) * mask.
         cfg.ratio == "token":    rho_{i,t} = exp(lr_{i,t})
         cfg.ratio == "sequence": s_i = exp( sum_t lr_{i,t} / |y_i| )       (GSPO)
                                  and rho_{i,t} = s_i for every token t of response i
                                  (s_i[:, None].expand_as(logp))
    2. The term, with A = adv[:, None] and band [lo, hi] = [1 - eps_low, 1 + eps_high]:
         cfg.clip == "none":  term = rho * A
         cfg.clip == "ppo":   term = min(rho * A, clamp(rho, lo, hi) * A)
         cfg.clip == "cispo": term = clamp(rho, lo, hi).detach() * A * logp
    3. Return term * mask.

    Check yourself: at logp == old_logp every rho is 1, and for every variant the
    gradient of term.sum() with respect to logp[i, t] is A_i * mask[i, t]. They differ only off-policy:
    PPO zeroes the gradient of tokens that left the band in the direction of A;
    CISPO never zeroes a gradient, it only caps the weight.
    """
    # TODO: ratio (token or sequence), then the three clip modes.
    raise NotImplementedError("Task 3: token_surrogate")


# ---------------------------------------------------------------------------
# Task 4: KL estimators (knob 5)
# ---------------------------------------------------------------------------
def kl_per_token(logp, ref_logp, estimator="k3"):
    """Per-token estimate of KL(pi_theta || pi_ref) at the sampled tokens. Returns (N, T).

    With d = ref_logp - logp = log(pi_ref / pi_theta) at the sampled token:
        "k1":  -d                       (unbiased, can be negative)
        "k2":  0.5 * d**2               (biased, low variance; REINFORCE++ uses it in the loss)
        "k3":  exp(d) - d - 1           (unbiased, always >= 0; GRPO uses it in the loss)
    Chapter 11 derives these; here you just need them as functions.
    """
    # TODO
    raise NotImplementedError("Task 4: kl_per_token")


# ---------------------------------------------------------------------------
# Task 5: the whole loss
# ---------------------------------------------------------------------------
def pg_loss(logp, old_logp, ref_logp, rewards, mask, group_ids, cfg):
    """The configurable critic-free policy-gradient LOSS (a scalar to minimize).

    1. KL in the reward (cfg.kl == "reward"): replace each reward by
           r_i - cfg.beta * sum_t (old_logp - ref_logp)_{i,t} * mask_{i,t}
       (a sequence-level k1 penalty, computed from the sampling policy, no grad).
    2. adv = group_advantages(rewards, group_ids, cfg.baseline, cfg.scale, cfg.adv_eps)
    3. per_token = token_surrogate(logp, old_logp, adv, mask, cfg)
    4. KL in the loss (cfg.kl == "loss"):
           per_token = per_token - cfg.beta * kl_per_token(logp, ref_logp, cfg.kl_estimator) * mask
    5. Return  -aggregate(per_token, mask, cfg.aggregation, cfg.max_len).
       (Minus, because optimizers minimize and we want to maximize the objective.)

    ref_logp may be None when cfg.kl == "none".
    """
    # TODO: five short steps, using Tasks 1-4.
    raise NotImplementedError("Task 5: pg_loss")