"""Tests for the PPO lab. Run:  python -m pytest -q   (or -k task1, -k task2, ...)"""

import importlib
import math
import os

import pytest
import torch

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

from cartpole import CartPole  # noqa: E402
from common import ActorCritic, collect_rollout  # noqa: E402

EPS = 0.2

# The six hand-computed cases from the chapter's "four cases" table, with eps = 0.2.
# (A, rho, objective min(rho*A, clip(rho)*A), d objective / d rho)
CASES = [
    (+2.0, 1.1, 2.2, 2.0),    # A > 0, inside the band: normal gradient
    (+2.0, 1.5, 2.4, 0.0),    # A > 0, rho > 1+eps: clipped, NO gradient
    (+2.0, 0.5, 1.0, 2.0),    # A > 0, rho < 1-eps: min keeps the unclipped term, gradient pulls rho up
    (-1.0, 0.9, -0.9, -1.0),  # A < 0, inside the band: normal gradient
    (-1.0, 0.6, -0.8, 0.0),   # A < 0, rho < 1-eps: clipped, NO gradient
    (-1.0, 1.4, -1.4, -1.0),  # A < 0, rho > 1+eps: min keeps the unclipped term, gradient pushes rho down
]


def _single(A, rho):
    logp_old = torch.tensor([math.log(0.5)])
    logp_new = torch.tensor([math.log(0.5 * rho)], requires_grad=True)
    adv = torch.tensor([A])
    return logp_new, logp_old, adv


# ---------------------------------------------------------------- Task 1
@pytest.mark.parametrize("A,rho,obj,slope", CASES)
def test_task1_clipped_loss_cases(A, rho, obj, slope):
    logp_new, logp_old, adv = _single(A, rho)
    loss = ppo.clipped_surrogate_loss(logp_new, logp_old, adv, EPS)
    assert loss.dim() == 0, "the loss must be a scalar tensor (use .mean())"
    assert loss.item() == pytest.approx(-obj, abs=1e-5), (
        f"A={A}, rho={rho}: expected objective {obj} (loss {-obj}), got loss {loss.item():.5f}. "
        "Did you forget the minus sign, or take max instead of min?")
    loss.backward()
    # d loss / d logp_new = -(d obj / d rho) * (d rho / d logp_new) = -slope * rho
    expected_grad = -slope * rho
    assert logp_new.grad.item() == pytest.approx(expected_grad, abs=1e-5), (
        f"A={A}, rho={rho}: expected d loss / d logp_new = {expected_grad}, got {logp_new.grad.item():.5f}. "
        "The gradient should be exactly zero only when (A>0 and rho>1+eps) or (A<0 and rho<1-eps).")


def test_task1_clipped_loss_batch_mean():
    logp_old = torch.full((len(CASES),), math.log(0.5))
    logp_new = torch.log(0.5 * torch.tensor([c[1] for c in CASES]))
    adv = torch.tensor([c[0] for c in CASES])
    expected = -sum(c[2] for c in CASES) / len(CASES)
    loss = ppo.clipped_surrogate_loss(logp_new, logp_old, adv, EPS)
    assert loss.item() == pytest.approx(expected, abs=1e-5), f"expected the batch MEAN {expected:.5f}, got {loss.item():.5f}"


def test_task1_at_old_policy_equals_vanilla_pg():
    """At theta = theta_old (rho = 1) the clipped loss has the vanilla policy-gradient gradient."""
    torch.manual_seed(0)
    logp_old = torch.log(torch.rand(8) * 0.9 + 0.05)
    logp_new = logp_old.clone().requires_grad_(True)
    adv = torch.randn(8)
    ppo.clipped_surrogate_loss(logp_new, logp_old, adv, EPS).backward()
    assert torch.allclose(logp_new.grad, -adv / 8, atol=1e-6), (
        "at rho = 1 the gradient w.r.t. logp_new should be -A/B (the REINFORCE gradient)")


# ---------------------------------------------------------------- Task 2
def test_task2_approx_kl_hand_value():
    # rho = [2, 0.5]: k3 = (2 - 1 - ln 2) + (0.5 - 1 + ln 2) = 0.5, mean 0.25
    logp_old = torch.log(torch.tensor([0.25, 0.5]))
    logp_new = torch.log(torch.tensor([0.5, 0.25]))
    kl = ppo.approx_kl(logp_new, logp_old)
    assert isinstance(kl, float), "approx_kl should return a Python float (use .item())"
    assert kl == pytest.approx(0.25, abs=1e-6), f"expected 0.25, got {kl}"


def test_task2_approx_kl_zero_and_nonnegative():
    lp = torch.log(torch.tensor([0.1, 0.7, 0.2]))
    assert ppo.approx_kl(lp, lp) == pytest.approx(0.0, abs=1e-7), "identical policies must give KL 0"
    torch.manual_seed(1)
    for _ in range(20):
        a, b = torch.randn(16), torch.randn(16)
        assert ppo.approx_kl(a, b) >= 0.0, "the k3 estimator is nonnegative for every sample"


def test_task2_approx_kl_is_unbiased_for_categorical():
    # Sample from pi_old, average k3 over many samples, compare with the exact KL(pi_old || pi_new).
    p_old = torch.tensor([0.5, 0.3, 0.2])
    p_new = torch.tensor([0.4, 0.4, 0.2])
    exact = (p_old * (p_old / p_new).log()).sum().item()
    g = torch.Generator().manual_seed(0)
    a = torch.multinomial(p_old, 200_000, replacement=True, generator=g)
    est = ppo.approx_kl(p_new.log()[a], p_old.log()[a])
    assert est == pytest.approx(exact, rel=0.05), f"exact KL {exact:.5f}, your estimate {est:.5f}"


def test_task2_clip_fraction():
    logp_old = torch.full((len(CASES),), math.log(0.5))
    logp_new = torch.log(0.5 * torch.tensor([c[1] for c in CASES]))
    frac = ppo.clip_fraction(logp_new, logp_old, EPS)
    assert isinstance(frac, float), "clip_fraction should return a Python float"
    # ratios 1.1, 1.5, 0.5, 0.9, 0.6, 1.4: four are outside [0.8, 1.2]
    assert frac == pytest.approx(4 / 6, abs=1e-6), f"expected 4/6, got {frac}"


# ---------------------------------------------------------------- Task 3
def test_task3_value_loss_unclipped():
    v = torch.tensor([1.0, 2.0])
    R = torch.tensor([0.0, 4.0])
    out = ppo.value_loss(v, R)
    assert out.item() == pytest.approx(1.25), f"0.5 * mean(1, 4) = 1.25, got {out.item()}"


def test_task3_value_loss_clipped():
    v = torch.tensor([1.0, 2.0], requires_grad=True)
    R = torch.tensor([0.0, 4.0])
    v_old = torch.tensor([1.5, 1.5])
    # V_clip = [1.3, 1.7]; squared errors: unclipped [1, 4], clipped [1.69, 5.29]; max -> [1.69, 5.29]
    out = ppo.value_loss(v, R, v_old, 0.2)
    assert out.item() == pytest.approx(0.5 * (1.69 + 5.29) / 2, abs=1e-5), f"expected 1.745, got {out.item()}"
    out.backward()
    # both samples take the clipped branch, which is constant in v once v - v_old is outside the band
    assert torch.allclose(v.grad, torch.zeros(2)), "outside the clip band the clipped value loss has zero gradient"


# ---------------------------------------------------------------- Task 4
def _fixed_batch(seed=0, n=256):
    torch.manual_seed(seed)
    env = CartPole(seed=seed)
    agent = ActorCritic()
    batch, _, _ = collect_rollout(env, agent, env.reset(), n)
    return agent, batch


def test_task4_update_runs_and_reports():
    agent, batch = _fixed_batch()
    cfg = ppo.PPOConfig(epochs=4, minibatch_size=64, anneal_lr=False)
    opt = torch.optim.Adam(agent.parameters(), lr=1e-3)
    before = [p.detach().clone() for p in agent.parameters()]
    stats = ppo.ppo_update(agent, opt, batch, cfg, generator=torch.Generator().manual_seed(0))
    for k in ["policy_loss", "value_loss", "entropy", "approx_kl", "clip_frac", "epochs_run"]:
        assert k in stats, f"stats is missing key {k!r}"
    assert stats["epochs_run"] == 4, f"with target_kl=None all 4 epochs should run, got {stats['epochs_run']}"
    assert 0.0 <= stats["clip_frac"] <= 1.0
    assert stats["approx_kl"] >= 0.0
    assert 0.0 < stats["entropy"] <= math.log(2) + 1e-6, "entropy of a 2-action policy is in (0, ln 2]"
    changed = any(not torch.equal(a, b.detach()) for a, b in zip(before, agent.parameters()))
    assert changed, "ppo_update should change the agent's parameters"


def test_task4_improves_surrogate_and_value_fit():
    agent, batch = _fixed_batch(seed=1)
    adv = batch["advantages"]
    adv_n = (adv - adv.mean()) / (adv.std() + 1e-8)

    def surrogate_and_vloss():
        with torch.no_grad():
            lp = agent.dist(batch["obs"]).log_prob(batch["actions"])
            s = (torch.exp(lp - batch["logp_old"]) * adv_n).mean().item()
            vl = ((agent.value(batch["obs"]) - batch["returns"]) ** 2).mean().item()
        return s, vl

    s0, v0 = surrogate_and_vloss()
    cfg = ppo.PPOConfig(epochs=5, minibatch_size=64, anneal_lr=False)
    ppo.ppo_update(agent, torch.optim.Adam(agent.parameters(), lr=1e-3), batch, cfg,
                   generator=torch.Generator().manual_seed(0))
    s1, v1 = surrogate_and_vloss()
    assert s1 > s0, f"the (unclipped) surrogate on the batch should go up: {s0:.4f} -> {s1:.4f}"
    assert v1 < v0, f"the critic's squared error on the batch should go down: {v0:.3f} -> {v1:.3f}"


def test_task4_early_stopping():
    agent, batch = _fixed_batch(seed=2)
    cfg = ppo.PPOConfig(epochs=10, minibatch_size=64, target_kl=1e-9, anneal_lr=False)
    stats = ppo.ppo_update(agent, torch.optim.Adam(agent.parameters(), lr=1e-3), batch, cfg,
                           generator=torch.Generator().manual_seed(0))
    assert stats["epochs_run"] == 1, (
        f"with a tiny target_kl the update should stop after the first epoch, ran {stats['epochs_run']}")


def test_task4_ppo_learns_cartpole():
    hist = ppo.train_ppo(ppo.PPOConfig(), total_steps=15_360, seed=0)
    final = hist["mean_return"][-1]
    assert final >= 150, (
        f"after 15k steps PPO should balance the pole for 150+ steps on average (reference gets ~400); got {final:.1f}")