"""Part 1: measuring reward overoptimization with best-of-n.

The setup is the toy model from the chapter: every response has a proxy score (what the
reward model says) and a gold score (what we really want). We draw a big pool of
responses from the base policy once, then ask: if we sampled n responses and kept the
one with the highest *proxy*, what proxy and gold scores would we get on average?

Fill in the three TODOs. Test with:  python -m pytest -q -k task1
"""

import numpy as np


def bon_kl(n):
    """Task 1a. KL(best-of-n policy || base policy) in nats: log n - (n - 1) / n.

    Exact when the proxy is continuous (no ties); an upper bound in general
    (Beirami et al. 2024). Must work elementwise on arrays of n as well as on ints.

    Examples: bon_kl(1) = 0.0, bon_kl(2) = log 2 - 1/2 ≈ 0.193.
    """
    # TODO: one line of numpy. Convert n to a float array first.
    raise NotImplementedError("Task 1a: bon_kl")


def best_of_n_stats(proxy, gold, n):
    """Task 1b. Expected proxy and gold score of best-of-n, computed EXACTLY from a pool.

    Draw n responses uniformly at random (with replacement) from the pool of M
    responses and keep the one with the highest proxy. Sort the pool by proxy,
    ascending, and give ranks k = 1..M. The winner has rank k with probability

        w_k = (k / M)^n - ((k - 1) / M)^n

    (all n draws have rank <= k, but not all have rank <= k - 1). So

        E[proxy of winner] = sum_k w_k * proxy_(k),   E[gold of winner] = sum_k w_k * gold_(k)

    where proxy_(k), gold_(k) belong to the response of rank k. No sampling needed:
    this is exact for the pool, and it is how Gao et al. computed their best-of-n curves.

    Args:
        proxy, gold: arrays of shape (M,), one entry per pooled response.
        n: int >= 1.
    Returns:
        (expected_proxy, expected_gold) as Python floats.
    """
    # TODO:
    #   1. order = np.argsort(proxy)          (ranks: ascending proxy)
    #   2. k = 1..M, the weights w_k above    (check: they sum to 1)
    #   3. weighted sums of proxy[order] and gold[order]
    raise NotImplementedError("Task 1b: best_of_n_stats")


def fit_gao_bon(d, gold_gain):
    """Task 1c. Least-squares fit of Gao et al.'s best-of-n form  gold_gain(d) = alpha*d - beta*d^2.

    Here d = sqrt(KL). The model is linear in (alpha, beta), so solve the 2-column
    least-squares problem  [d, -d^2] @ [alpha, beta] ≈ gold_gain  (no intercept: at
    d = 0 you are the base policy, so the gain is 0).

    Args:
        d, gold_gain: arrays of shape (K,).
    Returns:
        (alpha, beta) as Python floats. A positive beta means gold eventually falls.
    """
    # TODO: build X with columns d and -d**2, then np.linalg.lstsq(X, gold_gain, rcond=None).
    raise NotImplementedError("Task 1c: fit_gao_bon")