"""Lab 11: the language model as a policy.

Fill in the five tasks marked TODO, in order. Each docstring gives the inputs,
outputs, shapes, and the formula from Chapter 11 (/rl/learn/11-llm-as-policy/).

    Task 1a  response_logprobs        per-token log-probs of the response (the shift and the gather)
    Task 1b  sequence_logprob         sum them with a mask -> log pi(y | x)
    Task 2   sample_with_temperature  draw next tokens from softmax(logits / T)
    Task 3   kl_estimators            k1, k2, k3 from log pi and log pi_ref of sampled tokens
    Task 4   tilted_optimum           pi* = pi_ref * exp(r / beta) / Z over whole responses,
             kl_objective             and J(pi) = E_pi[r] - beta * KL(pi || pi_ref)
    Task 5   token_level_optimum      the same pi*, built token by token from soft values

`generate` at the bottom is complete: it runs the sampling loop with your
Task 2 function. Run the tests with `python -m pytest -q`.
"""

import numpy as np
import torch
import torch.nn.functional as F

from tinylm import BOS, EOS, PAD, encode_prompts


# ----------------------------------------------------------------------------- Task 1
def response_logprobs(logits, ids, prompt_len):
    """Log-probability of every response token under the model.

    The model reads ids[:, :L] and returns logits[:, t, :] = the distribution of
    the token at position t + 1 (position t predicts the NEXT token). The
    response starts at column `prompt_len`.

    Args:
        logits:     (B, L, V) float tensor, the model's output on `ids` (all columns).
        ids:        (B, L) long tensor, prompt (columns < prompt_len) then response.
        prompt_len: int, the column where the response starts.

    Returns:
        (B, L - prompt_len) tensor: entry [b, k] is
        log pi(ids[b, prompt_len + k] | ids[b, :prompt_len + k]).

    Hints: log_softmax over the vocabulary; logits at column t go with the token at
    column t + 1; torch.gather picks one entry per position. No mask here: padding
    positions get whatever value the model gives them (Task 1b removes them).
    """
    # TODO: implement (about 4 lines)
    raise NotImplementedError("Task 1a: response_logprobs")


def sequence_logprob(tok_logp, mask):
    """log pi(y | x) = sum over response tokens of log pi(y_t | x, y_<t).

    Args:
        tok_logp: (B, T) per-token log-probs (from response_logprobs). Positions
                  where mask == 0 may contain anything, including -inf or nan.
        mask:     (B, T) float or bool, 1 for real response tokens (incl. <eos>),
                  0 for padding after <eos>.

    Returns:
        (B,) tensor of sequence log-probs.

    Careful: 0 * (-inf) = nan in floating point, so multiplying by the mask is
    not enough. Use torch.where (or masked_fill) to replace masked entries by 0.
    """
    # TODO: implement (1-2 lines)
    raise NotImplementedError("Task 1b: sequence_logprob")


# ----------------------------------------------------------------------------- Task 2
def sample_with_temperature(logits, temperature, generator=None):
    """Sample one token per row from softmax(logits / temperature).

    Args:
        logits:      (B, V) float tensor (may contain -inf for forbidden tokens).
        temperature: float >= 0. temperature == 0 means greedy (argmax).
        generator:   optional torch.Generator, pass it to torch.multinomial so
                     results are reproducible.

    Returns:
        (B,) long tensor of token ids.

    Formula (Chapter 11, "Temperature"): pi_T(a) = exp(z_a / T) / sum_b exp(z_b / T).
    """
    # TODO: implement (about 4 lines)
    raise NotImplementedError("Task 2: sample_with_temperature")


# ----------------------------------------------------------------------------- Task 3
def kl_estimators(logp, ref_logp):
    """Per-token KL estimators for KL(pi || pi_ref) from tokens sampled from pi.

    Let l = logp - ref_logp = log pi(y_t | s_t) / pi_ref(y_t | s_t) for the sampled token.
        k1 = l
        k2 = l**2 / 2
        k3 = exp(-l) - 1 + l          (= pi_ref/pi - 1 - log(pi_ref/pi))

    Args:
        logp, ref_logp: tensors of the same shape (any shape).

    Returns:
        (k1, k2, k3), three tensors of that shape.
    """
    # TODO: implement (about 4 lines)
    raise NotImplementedError("Task 3: kl_estimators")


# ----------------------------------------------------------------------------- Task 4
def tilted_optimum(ref_probs, rewards, beta):
    """The KL-regularized optimum over an enumerable set of responses.

        pi*(y) = pi_ref(y) * exp(r(y) / beta) / Z,     Z = sum_y pi_ref(y) exp(r(y) / beta)

    Args:
        ref_probs: numpy array (any shape, e.g. (V, V) for the 2-token toy LM) of
                   reference probabilities, summing to 1. Some entries may be 0.
        rewards:   numpy array of the same shape.
        beta:      float > 0.

    Returns:
        (pi_star, log_Z): pi_star has the same shape as ref_probs and sums to 1;
        log_Z is a float.

    Compute in log space (subtract the max before exponentiating) so that a small
    beta does not overflow. Responses with ref_probs == 0 must get exactly 0.
    """
    # TODO: implement (about 6 lines)
    raise NotImplementedError("Task 4: tilted_optimum")


def kl_objective(probs, ref_probs, rewards, beta):
    """J(pi) = E_pi[r] - beta * KL(pi || pi_ref), with 0 * log 0 = 0. Returns a float."""
    # TODO: implement (about 4 lines)
    raise NotImplementedError("Task 4: kl_objective")


# ----------------------------------------------------------------------------- Task 5
def token_level_optimum(first, second, rewards, beta):
    """Build pi* one token at a time, for the 2-token toy LM, using soft values.

    Work backwards from the end of the response (Chapter 11, "The optimum, one
    token at a time"):
        V(i)        = beta * log sum_j second[i, j] * exp(rewards[i, j] / beta)
        V0          = beta * log sum_i first[i]     * exp(V(i) / beta)        (= beta log Z)
        pi*(j | i)  = second[i, j] * exp((rewards[i, j] - V(i)) / beta)
        pi*(i)      = first[i]     * exp((V(i) - V0) / beta)

    Args:
        first:   (V,) reference first-token probabilities pi_ref(y1 = i).
        second:  (V, V) reference second-token probabilities pi_ref(y2 = j | y1 = i).
        rewards: (V, V) reward of response "ij".
        beta:    float > 0.

    Returns:
        (first_star, second_star, V, V0): (V,), (V, V), (V,), float.
        first_star[:, None] * second_star must equal tilted_optimum's pi*.

    Use a stable log-sum-exp (e.g. np.logaddexp.reduce or subtract the max).
    """
    # TODO: implement (about 8 lines)
    raise NotImplementedError("Task 5: token_level_optimum")


# ----------------------------------------------------------------------------- provided
@torch.no_grad()
def generate(model, prompts, temperature=1.0, max_new=4, generator=None, sampler=None):
    """Autoregressive generation with your `sample_with_temperature` (complete code).

    The token-level MDP in action: state = prompt + tokens so far, action = next
    token, transition = append it. Returns the same dict as tinylm.sample:
        ids (B, L + max_new), prompt_len L, mask (B, max_new) with 1 on generated tokens.
    """
    sampler = sampler or sample_with_temperature
    model.eval()
    ids, L = encode_prompts(prompts)
    B = ids.shape[0]
    done = torch.zeros(B, dtype=torch.bool)
    mask = torch.zeros(B, max_new)
    for t in range(max_new):
        logits = model(ids)[:, -1, :].clone()
        logits[:, [BOS, PAD]] = float("-inf")  # never generate these
        nxt = sampler(logits, temperature, generator)
        nxt = torch.where(done, torch.full_like(nxt, PAD), nxt)
        mask[:, t] = (~done).float()
        ids = torch.cat([ids, nxt[:, None]], dim=1)
        done |= nxt == EOS
    model.train()
    return {"ids": ids, "prompt_len": L, "mask": mask}