"""Tests for the Chapter 7 lab. Run from this folder:

    python -m pytest -q                  # tests your pg.py
    RL_SOLUTIONS=1 python -m pytest -q   # tests the reference solution
"""

import importlib
import os

import numpy as np
import pytest
import torch

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

# The chapter's worked example: 3 arms with mean rewards 1, 2, 4, uniform policy.
Q = np.array([1.0, 2.0, 4.0])
THETA0 = np.zeros(3)


# ----------------------------------------------------------------- Task 1
def test_score_uniform_policy_by_hand():
    s = pg.softmax_score(THETA0, np.array([0, 2]))
    expected = np.array([[2 / 3, -1 / 3, -1 / 3],
                         [-1 / 3, -1 / 3, 2 / 3]])
    assert s.shape == (2, 3), f"expected shape (2, 3), got {np.shape(s)}"
    np.testing.assert_allclose(s, expected, atol=1e-12,
                               err_msg="score of a uniform 3-way softmax should be e_a - (1/3, 1/3, 1/3)")


def test_score_matches_finite_differences():
    theta = np.array([0.3, -1.2, 0.8, 2.0])
    eps = 1e-6
    for a in range(4):
        num = np.zeros(4)
        for k in range(4):
            d = np.zeros(4)
            d[k] = eps
            num[k] = (np.log(pg.softmax(theta + d)[a]) - np.log(pg.softmax(theta - d)[a])) / (2 * eps)
        got = pg.softmax_score(theta, np.array([a]))[0]
        np.testing.assert_allclose(got, num, atol=1e-6,
                                   err_msg=f"score for action {a} disagrees with a finite-difference gradient of log pi")


def test_score_has_zero_mean_under_pi():
    theta = np.array([0.5, -0.3, 1.1])
    pi = pg.softmax(theta)
    s = pg.softmax_score(theta, np.arange(3))
    np.testing.assert_allclose(pi @ s, np.zeros(3), atol=1e-12,
                               err_msg="sum_a pi(a) * score(a) must be the zero vector")


# ----------------------------------------------------------------- Task 2
def test_exact_gradient_worked_example():
    g = pg.exact_gradient(THETA0, Q)
    np.testing.assert_allclose(g, [-4 / 9, -1 / 9, 5 / 9], atol=1e-12,
                               err_msg="chapter's worked example: grad J at uniform policy should be (-4/9, -1/9, 5/9)")


def test_exact_gradient_matches_finite_differences():
    rng = np.random.default_rng(0)
    theta = rng.normal(size=10)
    q = rng.normal(size=10)
    J = lambda th: pg.softmax(th) @ q
    eps = 1e-6
    num = np.array([(J(theta + eps * e) - J(theta - eps * e)) / (2 * eps) for e in np.eye(10)])
    np.testing.assert_allclose(pg.exact_gradient(theta, q), num, atol=1e-7,
                               err_msg="exact_gradient disagrees with finite differences of J on a 10-armed bandit")


# ----------------------------------------------------------------- Task 3
def test_reinforce_estimates_by_hand():
    est = pg.reinforce_estimates(THETA0, np.array([2, 0]), np.array([4.0, 1.0]), baseline=1.0)
    expected = np.array([[-1.0, -1.0, 2.0],      # (4 - 1) * (-1/3, -1/3, 2/3)
                         [0.0, 0.0, 0.0]])       # (1 - 1) * anything
    np.testing.assert_allclose(est, expected, atol=1e-12,
                               err_msg="row i should be (R_i - b) * (e_{a_i} - pi)")


def test_exact_average_over_arms_equals_gradient():
    # With noiseless rewards, averaging the estimate over arms with weights pi gives grad J exactly,
    # for ANY baseline. This is the whole chapter in one assertion.
    theta = np.array([0.2, -0.7, 0.4])
    pi = pg.softmax(theta)
    for b in [0.0, 2.5, -10.0]:
        est = pg.reinforce_estimates(theta, np.arange(3), Q, baseline=b)
        np.testing.assert_allclose(pi @ est, pg.exact_gradient(theta, Q), atol=1e-12,
                                   err_msg=f"E[(R - b) * score] should equal grad J (baseline b={b})")


@pytest.mark.parametrize("baseline", [0.0, 7.0])
def test_monte_carlo_estimate_is_unbiased(baseline):
    rng = np.random.default_rng(1)
    theta = np.array([0.5, 0.0, -0.5])
    a, r = pg.sample_bandit(theta, Q, 200_000, rng, noise=1.0)
    est = pg.reinforce_estimates(theta, a, r, baseline=baseline).mean(axis=0)
    np.testing.assert_allclose(est, pg.exact_gradient(theta, Q), atol=0.03,
                               err_msg=f"the sample mean of REINFORCE estimates should match grad J (b={baseline})")


def test_baseline_reduces_variance():
    rng = np.random.default_rng(2)
    shifted = Q + 10.0  # all rewards large and positive
    a, r = pg.sample_bandit(THETA0, shifted, 50_000, rng, noise=1.0)
    J = pg.softmax(THETA0) @ shifted
    v0 = pg.total_variance(pg.reinforce_estimates(THETA0, a, r, baseline=0.0))
    vJ = pg.total_variance(pg.reinforce_estimates(THETA0, a, r, baseline=J))
    assert vJ < v0 / 20, f"baseline b = J should cut the variance a lot here: got {vJ:.3f} vs {v0:.3f} without"


# ----------------------------------------------------------------- Task 4
def test_rewards_to_go_by_hand():
    np.testing.assert_allclose(pg.rewards_to_go(np.array([1.0, 0.0, 2.0]), 0.5), [1.5, 1.0, 2.0],
                               err_msg="G_2 = 2, G_1 = 0 + 0.5*2 = 1, G_0 = 1 + 0.5*1 = 1.5")
    np.testing.assert_allclose(pg.rewards_to_go(np.ones(4), 1.0), [4, 3, 2, 1],
                               err_msg="with gamma = 1 and reward 1 per step, G_t counts the remaining steps")
    np.testing.assert_allclose(pg.rewards_to_go(np.array([0.0, 1.0]), 0.9), [0.9, 1.0])


def test_rewards_to_go_returns_floats_and_shape():
    out = pg.rewards_to_go(np.array([1, 1, 1]), 0.9)
    assert np.shape(out) == (3,), f"expected shape (3,), got {np.shape(out)}"
    np.testing.assert_allclose(out, [2.71, 1.9, 1.0], err_msg="integer rewards should still give float returns")


# ----------------------------------------------------------------- Task 5
def test_reinforce_loss_value():
    logits = torch.tensor([[0.0, 0.0], [np.log(3.0), 0.0]], requires_grad=True)
    actions = torch.tensor([1, 0])
    weights = torch.tensor([2.0, -1.0])
    loss = float(pg.reinforce_loss(logits, actions, weights).detach())
    # log pi = log(1/2) and log(3/4)
    expected = -(2.0 * np.log(0.5) + (-1.0) * np.log(0.75)) / 2
    assert abs(loss - expected) < 1e-6, f"expected loss {expected:.6f}, got {loss:.6f}"


def test_reinforce_loss_gradient_is_minus_weighted_score():
    torch.manual_seed(0)
    logits = torch.randn(5, 3, requires_grad=True)
    actions = torch.tensor([0, 2, 1, 1, 0])
    weights = torch.tensor([1.0, -2.0, 0.5, 3.0, 0.0])
    pg.reinforce_loss(logits, actions, weights).backward()
    pi = torch.softmax(logits.detach(), -1)
    onehot = torch.nn.functional.one_hot(actions, 3).float()
    expected = -(weights[:, None] * (onehot - pi)) / 5
    assert torch.allclose(logits.grad, expected, atol=1e-6), \
        "d loss / d logits should be -(1/N) * w_i * (e_{a_i} - pi_i)"


def test_reinforce_loss_does_not_backprop_into_weights():
    logits = torch.zeros(2, 2, requires_grad=True)
    weights = torch.tensor([1.0, 2.0], requires_grad=True)
    pg.reinforce_loss(logits, torch.tensor([0, 1]), weights).backward()
    assert weights.grad is None or torch.all(weights.grad == 0), \
        "weights are constants in the REINFORCE loss: detach them"


# ----------------------------------------------------------------- does it learn?
def test_reinforce_with_baseline_learns_cartpole():
    hist = pg.train_reinforce("baseline", n_iters=30, seed=1)
    late = float(np.mean(hist[-40:]))
    assert late > 150, f"REINFORCE with a baseline should reach a mean return > 150 in 30 iterations; got {late:.1f}"