"""Tests for gibbs.py. Run: python -m pytest -q test_gibbs.py"""

import importlib
import os

import numpy as np
import pytest

gibbs = importlib.import_module("solutions.gibbs" if os.environ.get("RL_SOLUTIONS") else "gibbs")

PI0 = np.array([0.5, 0.3, 0.2])
R = np.array([0.0, 1.0, 2.0])


def log_Z(pi0, r, beta):
    return np.log(np.sum(pi0 * np.exp(r / beta)))


# ------------------------------------------------------------ Task 5a
def test_worked_example_from_chapter():
    pi = gibbs.gibbs_policy(PI0, R, 1.0)
    # weights 0.5, 0.3e, 0.2e^2 = 0.5, 0.8155, 1.4778; Z = 2.7933
    assert np.allclose(pi, [0.1790, 0.2919, 0.5291], atol=1e-3), f"expected about [0.179, 0.292, 0.529], got {pi}"
    assert pi.sum() == pytest.approx(1.0)


def test_limits_of_beta():
    big = gibbs.gibbs_policy(PI0, R, 1e4)
    assert np.allclose(big, PI0, atol=1e-3), "as beta -> infinity the optimum should be the reference pi0"
    small = gibbs.gibbs_policy(PI0, R, 1e-3)
    assert np.allclose(small, [0, 0, 1], atol=1e-6), "as beta -> 0 the optimum should be greedy (all mass on the best action)"


def test_numerically_stable():
    pi = gibbs.gibbs_policy(PI0, np.array([1000.0, 1001.0, 999.0]), 0.01)
    assert np.all(np.isfinite(pi)), "no overflow: subtract the max logit before exponentiating"
    assert np.allclose(pi, [0, 1, 0], atol=1e-6)


def test_zero_reference_mass_stays_zero():
    pi = gibbs.gibbs_policy(np.array([0.0, 0.6, 0.4]), np.array([100.0, 0.0, 1.0]), 1.0)
    assert pi[0] == 0.0, "an action the reference never takes gets zero probability, whatever its reward"
    assert pi.sum() == pytest.approx(1.0)


def test_uniform_reference_is_softmax():
    r = np.array([0.1, 0.3, 0.6, 0.2])
    pi = gibbs.gibbs_policy(np.ones(4) / 4, r, 0.5)
    soft = np.exp(r / 0.5) / np.exp(r / 0.5).sum()
    assert np.allclose(pi, soft), "with a uniform reference, pi* is softmax(r / beta) (max-entropy RL)"


# ------------------------------------------------------------ Task 5b
def test_objective_worked_example():
    pi = np.array([0.1790, 0.2919, 0.5291])
    j = gibbs.kl_objective(pi, PI0, R, 1.0)
    assert j == pytest.approx(1.0272, abs=2e-3), f"J(pi*) in the worked example is about 1.027, got {j}"
    assert gibbs.kl_objective(PI0, PI0, R, 1.0) == pytest.approx(0.7), "J(pi0) = E_pi0[r] = 0.3 + 0.4 = 0.7 (KL is 0)"


def test_objective_handles_zeros():
    j = gibbs.kl_objective(np.array([0.0, 0.0, 1.0]), PI0, R, 1.0)
    assert j == pytest.approx(2.0 - np.log(1 / 0.2)), "greedy policy: J = 2 - log(1/0.2); 0 log 0 terms must count as 0"


# ------------------------------------------------------------ the variational principle
@pytest.mark.parametrize("beta", [0.1, 0.5, 1.0, 3.0])
def test_gibbs_policy_maximizes_objective(beta):
    rng = np.random.default_rng(0)
    pi0 = np.array([0.40, 0.25, 0.15, 0.10, 0.07, 0.03])
    r = np.array([0.1, 0.3, 0.6, 0.2, 0.9, 1.0])
    star = gibbs.gibbs_policy(pi0, r, beta)
    j_star = gibbs.kl_objective(star, pi0, r, beta)
    assert j_star == pytest.approx(beta * log_Z(pi0, r, beta), abs=1e-9), "the maximum value should equal beta * log Z"
    for _ in range(2000):
        p = rng.dirichlet(np.ones(6)) if rng.random() < 0.5 else star * np.exp(0.2 * rng.standard_normal(6))
        p = p / p.sum()
        j = gibbs.kl_objective(p, pi0, r, beta)
        assert j <= j_star + 1e-12, "found a policy that beats pi*: something is off"
        # the exact identity J(p) = beta log Z - beta KL(p || pi*)
        kl_p_star = np.sum(p * np.log(p / star))
        assert j == pytest.approx(beta * log_Z(pi0, r, beta) - beta * kl_p_star, abs=1e-9)