"""Tests for Parts 2-3 (failures.py). Run:  python -m pytest -q -k task2  (or task3, task4, smoke)"""

import importlib
import math
import os

import pytest
import torch

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

from common import RLConfig, proxy_reward, rl_train  # noqa: E402
from tinylm import TinyLM  # noqa: E402


# ---------------------------------------------------------------- Task 2
def test_task2_kl_penalized_rewards_hand_example():
    rewards = torch.tensor([1.0, 0.0])
    logp = torch.tensor([[-0.1, -0.2, -5.0], [-1.0, -1.0, -1.0]])
    ref = torch.tensor([[-0.5, -0.4, -9.0], [-1.0, -2.0, -1.0]])
    mask = torch.tensor([[1.0, 1.0, 0.0], [1.0, 1.0, 1.0]])
    # response 0: (0.4 + 0.2) = 0.6 (third token is padding, ignored); response 1: 0 + 1 + 0 = 1
    out = impl.kl_penalized_rewards(rewards, logp, ref, mask, beta=0.5)
    assert out.shape == (2,), f"expected one reward per response, shape (2,), got {tuple(out.shape)}"
    torch.testing.assert_close(out, torch.tensor([1.0 - 0.3, 0.0 - 0.5]),
                               msg="expected [0.7, -0.5]. Did you apply the mask before summing?")


def test_task2_beta_zero_is_identity():
    rewards = torch.tensor([0.3, -1.2, 2.0])
    logp, ref = torch.randn(3, 4), torch.randn(3, 4)
    out = impl.kl_penalized_rewards(rewards, logp, ref, torch.ones(3, 4), beta=0.0)
    torch.testing.assert_close(out, rewards)


# ---------------------------------------------------------------- Task 3
def _one(A, rho, eps_low, eps_high):
    old = torch.tensor([[math.log(0.5)]])
    new = torch.tensor([[math.log(0.5 * rho)]], requires_grad=True)
    loss = impl.clipped_pg_loss(new, old, torch.tensor([A]), torch.ones(1, 1), eps_low, eps_high)
    loss.backward()
    return loss.item(), new.grad.item()


@pytest.mark.parametrize("A,rho,el,eh,obj,slope", [
    (+1.0, 1.1, 0.2, 0.2, 1.1, 1.0),    # inside the band: normal gradient
    (+1.0, 1.25, 0.2, 0.2, 1.2, 0.0),   # symmetric clip: rho > 1.2, gradient switched off
    (+1.0, 1.25, 0.2, 0.28, 1.25, 1.0),  # clip-higher: 1.25 < 1.28, still pushing up
    (+1.0, 1.35, 0.2, 0.28, 1.28, 0.0),  # clip-higher, but above 1.28: off
    (-1.0, 0.7, 0.2, 0.28, -0.8, 0.0),  # negative advantage, below 1 - eps_low: off
    (-1.0, 1.4, 0.2, 0.28, -1.4, -1.0),  # negative advantage, rho too high: still pushes down
])
def test_task3_clip_cases(A, rho, el, eh, obj, slope):
    loss, grad = _one(A, rho, el, eh)
    assert loss == pytest.approx(-obj, abs=1e-5), (
        f"A={A}, rho={rho}, eps_low={el}, eps_high={eh}: expected objective {obj} (loss {-obj}), got loss {loss:.5f}")
    # d loss / d logp = -(d obj / d rho) * rho
    assert grad == pytest.approx(-slope * rho, abs=1e-5), (
        f"A={A}, rho={rho}: expected d loss/d logp = {-slope * rho:.4f}, got {grad:.4f}. "
        "Is the upper clip at 1 + eps_high and the lower at 1 - eps_low?")


def test_task3_token_mean_with_mask():
    logp = torch.zeros(2, 3)
    old = torch.zeros(2, 3)
    adv = torch.tensor([1.0, -2.0])
    mask = torch.tensor([[1.0, 1.0, 1.0], [1.0, 0.0, 0.0]])
    # rho = 1 everywhere: objective per token = A. Token mean = (1+1+1-2) / 4 = 0.25
    loss = impl.clipped_pg_loss(logp, old, adv, mask)
    assert loss.dim() == 0, "the loss must be a scalar"
    assert loss.item() == pytest.approx(-0.25), "expected the masked TOKEN mean: -(1+1+1-2)/4 = -0.25"


# ---------------------------------------------------------------- Task 4
def test_task4_entropy_known_values():
    V = 4
    uniform = torch.zeros(1, 1, V)
    peaked = torch.tensor([[[0.0, -1e9, -1e9, -1e9]]])
    logits = torch.cat([uniform, peaked], dim=1)  # (1, 2, V)
    H = impl.token_entropy(logits, torch.ones(1, 2))
    assert H.item() == pytest.approx(math.log(4) / 2, abs=1e-5), \
        "uniform over 4 has entropy log 4; a one-hot has 0; the mean is log(4)/2"
    assert not torch.isnan(H), "got nan: use log_softmax, not log(softmax), so 0 * log 0 = 0"
    H2 = impl.token_entropy(logits, torch.tensor([[1.0, 0.0]]))
    assert H2.item() == pytest.approx(math.log(4), abs=1e-5), "masked positions must not count"


def test_task4_entropy_two_point():
    p = 0.2
    logits = torch.tensor([[[math.log(p), math.log(1 - p)]]])
    H = impl.token_entropy(logits, torch.ones(1, 1))
    assert H.item() == pytest.approx(-(p * math.log(p) + (1 - p) * math.log(1 - p)), abs=1e-6)


# ---------------------------------------------------------------- smoke test of the full loop
def test_smoke_rl_loop_runs():
    torch.manual_seed(0)
    model = TinyLM()
    cfg = RLConfig(steps=2, prompts_per_step=4, group_size=4, beta=0.05, epochs=2, eps_high=0.28)
    hist = rl_train(model, proxy_reward, cfg, impl)
    for k in ("proxy", "gold", "kl", "entropy", "length"):
        v = hist[k]
        assert len(v) == 2 and all(math.isfinite(x) for x in v), f"history '{k}' should have 2 finite entries"
    assert hist["entropy"][0] > 0.5, "an untrained model should be close to uniform: high entropy"