"""Chapter 10 lab, part 2: the KL-regularized optimum (Gibbs variational principle).

Task 5a: gibbs_policy.  Task 5b: kl_objective.  Tests: test_gibbs.py.
numpy only.
"""

import numpy as np


# ---------------------------------------------------------------- Task 5a
def gibbs_policy(pi0, r, beta):
    """The KL-regularized optimal policy (the Gibbs / Boltzmann distribution)

        pi*(a) = pi0(a) * exp(r(a) / beta) / Z,   Z = sum_a pi0(a) exp(r(a) / beta).

    Parameters
    ----------
    pi0 : array (K,), reference distribution (non-negative, sums to 1; may contain zeros).
    r : array (K,), reward of each action.
    beta : float > 0, KL coefficient (temperature).

    Returns
    -------
    pi : array (K,), sums to 1. Must not overflow for r / beta as large as 1e5:
    work in log space (subtract the max before exponentiating).
    """
    # TODO: logits = log pi0 + r / beta; subtract the max; exponentiate; normalize.
    raise NotImplementedError("Task: implement gibbs_policy")


# ---------------------------------------------------------------- Task 5b
def kl_objective(pi, pi0, r, beta):
    """The KL-regularized objective

        J(pi) = sum_a pi(a) r(a)  -  beta * KL(pi || pi0),
        KL(pi || pi0) = sum_a pi(a) log( pi(a) / pi0(a) ),

    with the convention 0 * log(0 / q) = 0 (actions with pi(a) = 0 contribute
    nothing). Returns a float. (If pi puts mass where pi0 is 0, KL is +inf and
    J is -inf; you may ignore that case.)
    """
    # TODO: expected reward minus beta * KL, skipping actions with pi(a) == 0.
    raise NotImplementedError("Task: implement kl_objective")