"""Chapter 21 lab, Tasks 4-5: two small tools every frontier paper uses.

Task 4  pass_at_k          : the unbiased pass@k estimator (Chen et al. 2021), the
                             metric at the centre of the "does RL sharpen or expand?" debate.
Task 5  maxrl_advantages   : MaxRL's per-group advantage (Tajwar et al. 2026), a one-line
                             change to a REINFORCE-with-mean-baseline advantage.

numpy only. Tests: test_tools.py.
"""

import numpy as np  # noqa: F401


# ---------------------------------------------------------------- Task 4
def pass_at_k(n, c, k):
    """Unbiased estimate of pass@k from n samples of which c are correct.

    pass@k = probability that at least one of k samples (drawn without replacement
    from your n) is correct:

        pass@k = 1 - C(n - c, k) / C(n, k)

    Compute the ratio of binomials as a product so nothing overflows:

        C(n - c, k) / C(n, k) = prod_{i=0}^{k-1} (n - c - i) / (n - i)

    and return 1.0 when n - c < k (every k-subset must contain a correct sample).

    Parameters: ints with 0 <= c <= n and 1 <= k <= n (raise ValueError otherwise).
    Returns: a Python float in [0, 1].
    Example: n=4, c=2: pass@1 = 0.5, pass@2 = 1 - 1/6 = 5/6.
    """
    # TODO: validate arguments; handle n - c < k; return 1 - product.
    raise NotImplementedError("Task 4: implement pass_at_k")


# ---------------------------------------------------------------- Task 5
def maxrl_advantages(rewards):
    """MaxRL advantages for binary rewards, one group (or several) at a time.

    For a group of N rewards r_1..r_N in {0, 1} with mean p_hat:

        A_i = (r_i - p_hat) / p_hat      if p_hat > 0
        A_i = 0                          if p_hat = 0   (no success: no signal)

    Compare: the mean-baseline (RLOO / Dr. GRPO style) advantage is r_i - p_hat, and
    GRPO's is (r_i - p_hat) / std. MaxRL divides by the pass rate p_hat instead, which
    up-weights hard prompts (small p_hat) the most.

    Parameters
    ----------
    rewards : array-like of shape (N,) or (P, N): P groups of N rewards each.
              Groups are along the LAST axis.

    Returns
    -------
    numpy float array with the same shape.
    Example: [1, 0, 0, 0] -> p_hat = 0.25 -> [3, -1, -1, -1].
    """
    # TODO: group mean along the last axis (keepdims=True); divide where mean > 0.
    raise NotImplementedError("Task 5: implement maxrl_advantages")