"""Given code for Task 3: a tiny "language model" whose sampler and trainer disagree.

You do not need to edit this file. Read `make_toy` once.

The toy
-------
* Responses are sequences of T tokens from a vocabulary of V tokens. Small enough
  (V**T sequences, 4**6 = 4096 by default) that every expectation is an exact sum.
* The TRAINER's policy: at position t, pi(y_t = v) = softmax(theta[t])_v.
  (Tokens are independent across positions to keep the toy enumerable. A real LM
  conditions on the prefix; none of the estimators below care about that.)
* The SAMPLER (the "inference engine") computes the same logits with an error:
  mu(y_t = v) = softmax(theta[t] + mismatch * b[t])_v, where b is a fixed, heavy-tailed
  random table. Think of b as what different kernels, a lower-precision LM head or a
  quantized rollout model do to the logits. mismatch = 0 means sampler == trainer.
* Reward r(y) = 1 if at least half of the tokens are token 0, else 0.

Gradients are taken with respect to theta (shape (T, V)). For a softmax,
    d log pi(y_t) / d theta[t] = onehot(y_t) - softmax(theta[t]),
and it is zero with respect to every other row theta[s], s != t. We call the
(T, V) array with row t equal to onehot(y_t) - pi_t the "score" of token t.
"""

import itertools

import numpy as np


def log_softmax(z):
    z = z - z.max(axis=-1, keepdims=True)
    return z - np.log(np.exp(z).sum(axis=-1, keepdims=True))


def make_toy(T=6, V=4, mismatch=0.3, seed=0):
    """Enumerate everything about the toy. Returns a dict with:

      Y            (M, T)    every possible sequence, M = V**T
      logp_train   (M, T)    log pi(y_t) for each token of each sequence (trainer)
      logp_sampler (M, T)    log mu(y_t) (sampler)
      p_train      (M,)      pi(y) = prod_t pi(y_t)
      p_sampler    (M,)      mu(y)
      R            (M,)      reward of each sequence
      score        (M, T, V) score[m, t] = onehot(y_t) - pi_t  (gradient of log pi(y_t) w.r.t. theta[t])
      true_grad    (T, V)    exact gradient of J(theta) = E_{y ~ pi}[r(y)]
      theta, b, mismatch, T, V
    """
    rng = np.random.default_rng(seed)
    theta = rng.normal(0.0, 0.7, size=(T, V))
    b = rng.standard_t(3, size=(T, V))
    lp = log_softmax(theta)                       # (T, V)
    lm = log_softmax(theta + mismatch * b)        # (T, V)
    pi = np.exp(lp)

    Y = np.array(list(itertools.product(range(V), repeat=T)))  # (M, T)
    pos = np.arange(T)
    logp_train = lp[pos, Y]
    logp_sampler = lm[pos, Y]
    p_train = np.exp(logp_train.sum(1))
    p_sampler = np.exp(logp_sampler.sum(1))
    R = (np.mean(Y == 0, axis=1) >= 0.5).astype(float)
    score = np.eye(V)[Y] - pi[None]                # (M, T, V)
    # grad J = E_pi[ r(y) * sum_t grad log pi(y_t) ]   (REINFORCE, exact)
    true_grad = np.einsum("m,m,mtv->tv", p_train, R, score)
    return dict(Y=Y, logp_train=logp_train, logp_sampler=logp_sampler, p_train=p_train,
                p_sampler=p_sampler, R=R, score=score, true_grad=true_grad,
                theta=theta, b=b, mismatch=mismatch, T=T, V=V)


def sample_from_sampler(toy, n, rng):
    """Draw n sequences from the SAMPLER mu, as the inference engine would. Returns indices into toy['Y']."""
    p = toy["p_sampler"] / toy["p_sampler"].sum()
    return rng.choice(len(p), size=n, p=p)


def estimator_moments(toy, per_seq_estimate):
    """Exact mean, bias and variance of a single-sample gradient estimator.

    per_seq_estimate: (M, T, V) array, the estimate you would compute if the sampler
    produced sequence m. Since sequences arrive with probability mu(y_m), the estimator's
    mean is sum_m mu(y_m) * estimate[m].

    Returns dict(mean=(T, V), bias=float, var=float, rel_bias=float, rel_rmse=function):
      bias     = || mean - true_grad ||
      var      = E_mu || estimate - mean ||^2   (total variance of ONE sample)
      rel_rmse(N) = sqrt(bias^2 + var / N) / || true_grad ||, the relative error of a
                    batch average over N samples.
    """
    mu = toy["p_sampler"]
    mean = np.einsum("m,mtv->tv", mu, per_seq_estimate)
    bias = float(np.linalg.norm(mean - toy["true_grad"]))
    var = float(np.einsum("m,m->", mu, ((per_seq_estimate - mean) ** 2).sum(axis=(1, 2))))
    g = float(np.linalg.norm(toy["true_grad"]))
    return dict(mean=mean, bias=bias, var=var, rel_bias=bias / g,
                rel_rmse=lambda N: float(np.sqrt(bias ** 2 + var / N) / g))