"""An enumerable toy language model (complete helper code; you don't edit this).

Every response is exactly two tokens long, drawn from a tiny vocabulary, so the
whole response distribution fits in a small matrix and every quantity in
Chapter 11 can be computed exactly, with no sampling:

    first[i]       = pi(y1 = i | x)                      shape (V,)
    second[i, j]   = pi(y2 = j | x, y1 = i)               shape (V, V)
    P[i, j]        = first[i] * second[i, j]              probability of response "ij"
    reward[i, j]   = r(x, "ij")

The scenarios are the same as in the chapter's interactive demo.
"""

import numpy as np

SCENARIOS = {
    # The worked example in the chapter: tokens a, b.
    "worked": dict(
        tokens=["a", "b"],
        first=np.array([0.7, 0.3]),
        second=np.array([[0.6, 0.4],
                         [0.5, 0.5]]),
        reward=np.array([[0.0, 1.0],
                         [0.0, 2.0]]),
    ),
    # The most likely response is also the best one.
    "likely": dict(
        tokens=["a", "b", "c"],
        first=np.array([0.6, 0.3, 0.1]),
        second=np.array([[0.7, 0.2, 0.1],
                         [0.3, 0.5, 0.2],
                         [0.2, 0.3, 0.5]]),
        reward=np.array([[1.0, 0.3, 0.0],
                         [0.5, 0.2, 0.0],
                         [0.0, 0.0, 0.2]]),
    ),
    # A learned reward model with a blind spot: it scores the rare response
    # "cc" at 1.6, but its true value is 0.
    "hack": dict(
        tokens=["a", "b", "c"],
        first=np.array([0.5, 0.4, 0.1]),
        second=np.array([[0.6, 0.3, 0.1],
                         [0.2, 0.6, 0.2],
                         [0.4, 0.4, 0.2]]),
        reward=np.array([[0.2, 0.6, 0.0],
                         [0.3, 1.0, 0.4],
                         [0.1, 0.2, 1.6]]),
        true_reward=np.array([[0.2, 0.6, 0.0],
                              [0.3, 1.0, 0.4],
                              [0.1, 0.2, 0.0]]),
    ),
}


def seq_probs(first, second):
    """P[i, j] = first[i] * second[i, j]: the chain rule for a 2-token response."""
    return first[:, None] * second


def temperature_policy(first, second, T):
    """Token-level temperature: every next-token distribution p becomes p^(1/T) / sum."""
    def sharpen(p):
        with np.errstate(divide="ignore"):
            logits = np.log(p) / T
        logits = logits - logits.max(axis=-1, keepdims=True)
        w = np.exp(logits)
        return w / w.sum(axis=-1, keepdims=True)
    return sharpen(first), sharpen(second)


def expectation(P, r):
    return float(np.sum(P * r))


def kl(P, Q):
    """KL(P || Q) for two arrays of probabilities of the same shape, with 0 log 0 = 0."""
    P, Q = np.asarray(P, float), np.asarray(Q, float)
    m = P > 0
    return float(np.sum(P[m] * np.log(P[m] / Q[m])))


def random_toy_lm(V, rng):
    """A random 2-token LM with vocabulary size V and random rewards (for tests)."""
    first = rng.dirichlet(np.ones(V))
    second = rng.dirichlet(np.ones(V), size=V)
    reward = rng.normal(size=(V, V))
    return first, second, reward