"""Chapter 20 lab, Task 3: train-inference mismatch and truncated importance sampling.

The inference engine samples from mu; the trainer's policy is pi. They SHOULD be
the same distribution, but different kernels and precisions make them differ a
little. You will measure how much that biases the policy gradient, and fix it.

    Task 3a  importance_weights   per-token weights for five correction schemes
    Task 3b  pg_estimate          the per-sample REINFORCE gradient estimate, with weights

Run `python -m pytest -q -k task3`. The toy problem and the exact bias/variance
calculator are given in toy.py.

Shapes: N sampled sequences, T tokens each, vocabulary V.
"""

import numpy as np

MODES = ("none", "token", "token_tis", "seq", "seq_tis")


# ---------------------------------------------------------------------------
# Task 3a
# ---------------------------------------------------------------------------
def importance_weights(logp_train, logp_sampler, mode, C=2.0):
    """Per-token importance weights, shape (N, T).

    logp_train:   (N, T) log pi(y_t) computed by the TRAINER for each sampled token
    logp_sampler: (N, T) log mu(y_t) reported by the SAMPLER for the same tokens
    mode:         one of MODES (raise ValueError otherwise)
    C:            truncation level for the *_tis modes

    With w_t = pi(y_t) / mu(y_t) (a per-token ratio) and W = prod_t w_t = pi(y) / mu(y)
    (the whole-sequence ratio):

      'none'       1 everywhere: pretend the sampler was the trainer (what a naive
                   "on-policy" implementation does)
      'token'      w_t                       per-token importance sampling
      'token_tis'  min(w_t, C)               per-token TRUNCATED importance sampling
                                             (Yao et al. 2025; the TIS in OlmoRL)
      'seq'        W, the same for every token of the sequence   (exact, unbiased IS)
      'seq_tis'    min(W, C), the same for every token of the sequence

    Work in log space (the per-token log-ratio is logp_train - logp_sampler; the
    sequence log-ratio is its sum over t) and exponentiate at the end.
    Return a NEW array of shape (N, T) in every mode (np.broadcast_to(...).copy()
    is handy for the sequence modes).
    """
    # TODO
    raise NotImplementedError("Task 3a: importance_weights")


# ---------------------------------------------------------------------------
# Task 3b
# ---------------------------------------------------------------------------
def pg_estimate(rewards, score, weights, baseline=0.0):
    """Per-sample policy-gradient estimates, shape (N, T, V).

    rewards:  (N,)       r(y) of each sampled sequence
    score:    (N, T, V)  score[n, t] = gradient of log pi(y_t) w.r.t. theta[t]
                         (see toy.py: onehot(y_t) - pi_t)
    weights:  (N, T)     from importance_weights
    baseline: a constant subtracted from the reward

    The estimate for sample n is the (T, V) array whose row t is

        (r(y_n) - baseline) * weights[n, t] * score[n, t]

    i.e. REINFORCE, with each token's term multiplied by its weight. (Summed over
    t this is "sum_t w_t * A * grad log pi(y_t)", the chapter's corrected gradient.
    Each token's gradient only touches its own row theta[t], which is why the result
    keeps a separate row per t.) Average over n to get a batch estimate.
    """
    # TODO: one line with broadcasting
    raise NotImplementedError("Task 3b: pg_estimate")