"""Tests for the Chapter 8 lab. Run:  python -m pytest -q
Check the reference solution with:  RL_SOLUTIONS=1 python -m pytest -q
"""

import importlib
import math
import os

import numpy as np
import pytest
import torch

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

# The worked example from the chapter: 4 steps, reward only at the end, gamma = 1.
R = np.array([0.0, 0.0, 0.0, 1.0])
V = np.array([0.5, 0.4, 0.8, 0.6])
D = np.array([0, 0, 0, 1])


# ---------------------------------------------------------------- Task 1
def test_td_errors_chapter_example():
    d = ac.td_errors(R, V, 0.0, D, gamma=1.0)
    assert np.allclose(d, [-0.1, 0.4, -0.2, 0.4]), f"expected deltas [-0.1, 0.4, -0.2, 0.4], got {d}"


def test_td_errors_bootstraps_from_last_value():
    # Not done at the end: the last step bootstraps from last_value.
    d = ac.td_errors([1.0, 1.0], [2.0, 3.0], 4.0, [0, 0], gamma=0.5)
    # delta_0 = 1 + 0.5*3 - 2 = 0.5 ; delta_1 = 1 + 0.5*4 - 3 = 0
    assert np.allclose(d, [0.5, 0.0]), f"expected [0.5, 0.0], got {d}"


def test_td_errors_respects_done():
    # dones[0] = 1: s_1 is terminal, so values[1] (next episode's first state) is not used.
    d = ac.td_errors([1.0, 1.0], [2.0, 3.0], 4.0, [1, 0], gamma=0.5)
    assert np.allclose(d, [-1.0, 0.0]), f"expected [-1.0, 0.0] (no bootstrap after done), got {d}"


# ---------------------------------------------------------------- Task 2
def test_gae_chapter_example_lambda_half():
    adv, ret = ac.compute_gae(R, V, 0.0, D, gamma=1.0, lam=0.5)
    assert np.allclose(adv, [0.1, 0.4, 0.0, 0.4]), f"expected GAE(1, 0.5) = [0.1, 0.4, 0.0, 0.4], got {adv}"
    assert np.allclose(ret, [0.6, 0.8, 0.8, 1.0]), f"returns should be advantages + values, got {ret}"


def test_gae_lambda_zero_is_td_error():
    adv, _ = ac.compute_gae(R, V, 0.0, D, gamma=1.0, lam=0.0)
    assert np.allclose(adv, [-0.1, 0.4, -0.2, 0.4]), f"lambda = 0 should give the TD errors, got {adv}"


def test_gae_lambda_one_is_monte_carlo():
    adv, _ = ac.compute_gae(R, V, 0.0, D, gamma=1.0, lam=1.0)
    assert np.allclose(adv, [0.5, 0.6, 0.2, 0.4]), f"lambda = 1 should give G_t - V(s_t) = [0.5, 0.6, 0.2, 0.4], got {adv}"


def test_gae_lambda_one_matches_discounted_return_random():
    rs = np.random.default_rng(0)
    T, gamma = 9, 0.9
    r, v = rs.normal(size=T), rs.normal(size=T)
    dones = np.zeros(T)
    dones[-1] = 1
    G = np.zeros(T)
    run = 0.0
    for t in reversed(range(T)):
        run = r[t] + gamma * run
        G[t] = run
    adv, _ = ac.compute_gae(r, v, 123.0, dones, gamma=gamma, lam=1.0)
    assert np.allclose(adv, G - v), "with lambda = 1 and a terminated episode, GAE must equal G_t - V(s_t)"


def test_gae_equals_explicit_sum():
    rs = np.random.default_rng(1)
    T, gamma, lam = 7, 0.95, 0.8
    r, v, last = rs.normal(size=T), rs.normal(size=T), 0.7
    dones = np.zeros(T)
    delta = r + gamma * np.append(v[1:], last) - v
    expected = np.array([sum((gamma * lam) ** l * delta[t + l] for l in range(T - t)) for t in range(T)])
    adv, _ = ac.compute_gae(r, v, last, dones, gamma=gamma, lam=lam)
    assert np.allclose(adv, expected), "GAE should equal sum_l (gamma*lam)^l delta_{t+l}"


def test_gae_episode_boundary():
    # Two episodes in one rollout: steps 0-1 (ends at step 1), then steps 2-3 (cut off, bootstrap).
    r = [1.0, 1.0, 1.0, 1.0]
    v = [2.0, 1.0, 3.0, 2.0]
    dones = [0, 1, 0, 0]
    adv, _ = ac.compute_gae(r, v, 5.0, dones, gamma=0.5, lam=0.5)
    # deltas = [-0.5, 0, -1, 1.5]; A3 = 1.5, A2 = -1 + 0.25*1.5, A1 = 0 (no carry across the boundary), A0 = -0.5
    assert np.allclose(adv, [-0.5, 0.0, -0.625, 1.5]), f"expected [-0.5, 0, -0.625, 1.5], got {adv}"


def test_gae_batched_columns_are_independent():
    rs = np.random.default_rng(2)
    T, N = 6, 3
    r, v = rs.normal(size=(T, N)), rs.normal(size=(T, N))
    dones = (rs.random((T, N)) < 0.3).astype(float)
    last = rs.normal(size=N)
    adv, ret = ac.compute_gae(r, v, last, dones, gamma=0.9, lam=0.7)
    assert adv.shape == (T, N) and ret.shape == (T, N), f"expected shape {(T, N)}, got {adv.shape}"
    for j in range(N):
        a_j, _ = ac.compute_gae(r[:, j], v[:, j], last[j], dones[:, j], gamma=0.9, lam=0.7)
        assert np.allclose(adv[:, j], a_j), f"column {j} of a (T, N) rollout should match running it alone"


# ---------------------------------------------------------------- Task 3
def test_normalize_advantages():
    out = ac.normalize_advantages(torch.tensor([1.0, 2.0, 3.0, 4.0]))
    s = math.sqrt(1.25)
    expected = torch.tensor([-1.5, -0.5, 0.5, 1.5]) / s
    assert torch.allclose(out, expected, atol=1e-5), f"expected {expected} (population std), got {out}"


def test_normalize_constant_batch_is_finite():
    out = ac.normalize_advantages(torch.full((5,), 3.0))
    assert torch.all(torch.isfinite(out)) and torch.allclose(out, torch.zeros(5)), "a constant batch should map to zeros"


# ---------------------------------------------------------------- Task 4
def _loss_inputs():
    logits = torch.tensor([[0.0, 0.0], [math.log(3.0), 0.0]], requires_grad=True)  # pi = (.5,.5), (.75,.25)
    values = torch.tensor([1.0, 2.0], requires_grad=True)
    actions = torch.tensor([0, 1])
    advantages = torch.tensor([2.0, 1.0], requires_grad=True)
    returns = torch.tensor([2.0, 0.0], requires_grad=True)
    return logits, values, actions, advantages, returns


def test_actor_critic_loss_values():
    logits, values, actions, adv, ret = _loss_inputs()
    total, info = ac.actor_critic_loss(logits, values, actions, adv, ret, value_coef=0.5, entropy_coef=0.01)
    pl = -(math.log(0.5) * 2 + math.log(0.25) * 1) / 2  # 1.386294
    vl = (1.0 + 4.0) / 2  # 2.5
    h = (math.log(2) + -(0.75 * math.log(0.75) + 0.25 * math.log(0.25))) / 2  # 0.627741
    assert info["policy_loss"] == pytest.approx(pl, abs=1e-5), f"policy_loss: expected {pl:.6f}, got {info['policy_loss']}"
    assert info["value_loss"] == pytest.approx(vl, abs=1e-5), f"value_loss: expected {vl}, got {info['value_loss']}"
    assert info["entropy"] == pytest.approx(h, abs=1e-5), f"entropy: expected {h:.6f}, got {info['entropy']}"
    expected = pl + 0.5 * vl - 0.01 * h
    assert total.item() == pytest.approx(expected, abs=1e-5), f"total: expected {expected:.6f}, got {total.item()}"
    assert isinstance(info["policy_loss"], float), "info values should be python floats (use .item())"


def test_actor_critic_loss_gradients():
    logits, values, actions, adv, ret = _loss_inputs()
    total, _ = ac.actor_critic_loss(logits, values, actions, adv, ret, value_coef=0.5, entropy_coef=0.0)
    total.backward()
    assert adv.grad is None and ret.grad is None, "advantages and returns are targets: detach them"
    # d/dV of 0.5 * mean((V - R)^2) = (V - R) / B
    assert torch.allclose(values.grad, torch.tensor([-0.5, 1.0])), f"value gradient should be (V - R)/B, got {values.grad}"
    # policy part: -(1/B) * A * (onehot(a) - pi)
    expected = torch.tensor([[-0.5, 0.5], [0.375, -0.375]])
    assert torch.allclose(logits.grad, expected, atol=1e-5), f"logit gradient: expected {expected}, got {logits.grad}"


# ---------------------------------------------------------------- scaffolding + integration
def test_vec_env_flags_truncation_not_termination():
    envs = ac.VecCartPole(2, seed=0, max_steps=5)
    envs.reset()
    for k in range(5):
        _, _, term, trunc, _ = envs.step(np.array([k % 2, k % 2]))
    assert trunc.all() and not term.any(), "hitting the time limit is truncation, not termination"


def test_a2c_learns_cartpole():
    hist = ac.train_a2c(total_steps=80_000, log_every=10_000, seed=0)
    best = max(hist["avg_return"])
    assert best >= 150, f"A2C should reach a 20-episode average return of 150 within 80k steps (random policy: ~22); best was {best:.1f}"