"""Tests for Lab 12. Run:  python -m pytest -q   (or -k task1, -k task2, ...)"""

import importlib
import math
import os
import random

import pytest
import torch

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

from prefs import EOS, PAD, DigitRM, gold_score, labeler_rank  # noqa: E402
from tinylm import load_or_train_sft  # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))


# ---------------------------------------------------------------- Task 1: Bradley-Terry loss
@pytest.mark.parametrize("rw,rl,expected", [
    (0.0, 0.0, math.log(2)),                 # RM can't tell them apart: loss ln 2
    (2.0, 0.0, math.log(1 + math.exp(-2))),  # right, and confident: small loss
    (0.0, 2.0, math.log(1 + math.exp(2))),   # wrong, and confident: big loss
])
def test_task1_bt_loss_values(rw, rl, expected):
    loss = rlhf.bt_loss(torch.tensor([rw]), torch.tensor([rl]))
    assert loss.dim() == 0, "bt_loss must return a scalar (take the mean over pairs)"
    assert loss.item() == pytest.approx(expected, abs=1e-5), (
        f"r_w={rw}, r_l={rl}: expected -log sigmoid({rw - rl}) = {expected:.5f}, got {loss.item():.5f}")


def test_task1_bt_loss_mean_and_gradient():
    rw = torch.tensor([0.0, 2.0, 0.0], requires_grad=True)
    rl = torch.tensor([0.0, 0.0, 2.0])
    loss = rlhf.bt_loss(rw, rl)
    expected = (math.log(2) + math.log(1 + math.exp(-2)) + math.log(1 + math.exp(2))) / 3
    assert loss.item() == pytest.approx(expected, abs=1e-5), f"expected the MEAN over pairs, {expected:.5f}"
    loss.backward()
    # d loss / d r_w = -(1 - sigmoid(r_w - r_l)) / N
    want = [-(1 - 1 / (1 + math.exp(-d))) / 3 for d in (0.0, 2.0, -2.0)]
    assert rw.grad.tolist() == pytest.approx(want, abs=1e-5), (
        f"gradient w.r.t. r_chosen should be -(1 - sigmoid(delta))/N = {want}, got {rw.grad.tolist()}")


def test_task1_bt_loss_shift_invariant_and_stable():
    rw, rl = torch.tensor([0.3, -1.2]), torch.tensor([0.1, 0.4])
    a = rlhf.bt_loss(rw, rl).item()
    b = rlhf.bt_loss(rw + 7.0, rl + 7.0).item()
    assert a == pytest.approx(b, abs=1e-6), "adding the same constant to both rewards must not change the loss"
    big = rlhf.bt_loss(torch.tensor([-100.0]), torch.tensor([0.0])).item()
    assert math.isfinite(big) and big == pytest.approx(100.0, abs=1e-3), (
        "for delta = -100 the loss is about 100; got %r. Use F.logsigmoid, not log(sigmoid(.))" % big)


# ---------------------------------------------------------------- Task 2: rankings -> pairs
def test_task2_ranking_to_pairs():
    assert rlhf.ranking_to_pairs([2, 0, 1]) == [(2, 0), (2, 1), (0, 1)]
    assert rlhf.ranking_to_pairs([7]) == [], "a ranking of one item has no pairs"
    for K in (2, 4, 9):
        pairs = rlhf.ranking_to_pairs(list(range(K)))
        assert len(pairs) == K * (K - 1) // 2, f"K={K}: expected K choose 2 = {K * (K - 1) // 2} pairs"
        assert all(w < l for w, l in pairs), "each pair must be (higher-ranked, lower-ranked)"


def _synthetic_rankings(n_prompts, K=4, seed=0):
    """Well-formed answers ranked by the hidden labeler. Like samples from the SFT model, the K answers
    to one prompt share their leading digits and differ mostly in the last one."""
    gen = random.Random(seed)
    data = []
    for _ in range(n_prompts):
        ids = torch.full((K, 4), PAD, dtype=torch.long)
        mask = torch.zeros(K, 4)
        prefix = [gen.randint(1, 9) for _ in range(gen.randint(0, 2))]
        for k in range(K):
            toks = prefix + [gen.randint(0, 9), EOS]
            ids[k, : len(toks)] = torch.tensor(toks)
            mask[k, : len(toks)] = 1.0
        out = {"ids": ids, "mask": mask, "prompt_len": 0}
        g = gold_score(out)
        data.append({"out": out, "ranking": labeler_rank(g.tolist(), gen), "gold": g})
    return data


def test_task2_reward_model_learns_the_preference():
    """Uses your bt_loss and ranking_to_pairs inside the given train_reward_model."""
    train, test = _synthetic_rankings(400, seed=1), _synthetic_rankings(200, seed=2)
    rm = rlhf.train_reward_model(train)
    w = rm.w.detach()
    even, odd = w[0::2].mean().item(), w[1::2].mean().item()
    assert even > odd + 0.3, f"the RM should learn 'even digits good': mean even weight {even:.2f}, odd {odd:.2f}"
    acc, ceiling = rlhf.rm_accuracy(rm, test)
    assert acc > ceiling - 0.04, f"held-out RM accuracy {acc:.3f} should be close to the ceiling {ceiling:.3f}"


# ---------------------------------------------------------------- Task 3: per-token rewards
def test_task3_kl_penalized_rewards():
    logp = torch.tensor([[-1.0, -2.0, -0.5, -3.0],
                         [-0.3, -9.0, -9.0, -9.0],
                         [-0.2, -0.4, -0.6, -0.8]])
    ref = torch.tensor([[-1.5, -2.0, -1.0, -0.2],
                        [-0.1, -1.0, -1.0, -1.0],
                        [-0.2, -0.2, -0.2, -0.2]])
    mask = torch.tensor([[1.0, 1, 1, 0], [1, 0, 0, 0], [1, 1, 1, 1]])
    score = torch.tensor([2.0, -1.0, 0.5])
    want = torch.tensor([[-0.05, 0.0, 1.95, 0.0],
                         [-0.98, 0.0, 0.0, 0.0],
                         [0.0, 0.02, 0.04, 0.56]])
    got = rlhf.kl_penalized_rewards(logp, ref, score, mask, beta=0.1)
    assert got.shape == (3, 4)
    assert torch.allclose(got, want, atol=1e-6), (
        f"expected\n{want}\ngot\n{got}\nKL part -beta*(logp-ref) at real tokens, + score at the LAST real token, 0 at padding")
    total = score - 0.1 * ((logp - ref) * mask).sum(1)
    assert torch.allclose(got.sum(1), total, atol=1e-6), "each row should sum to score - beta * sequence log-ratio"


def test_task3_chapter_paris_example():
    """The chapter's worked example: "Paris", ".", EOS with r_phi = 1.5 and beta = 0.1 (plus one padding slot)."""
    logp = torch.tensor([[-0.3, -0.9, -0.1, -5.0]])
    ref = torch.tensor([[-0.8, -0.7, -0.1, -1.0]])
    mask = torch.tensor([[1.0, 1.0, 1.0, 0.0]])
    got = rlhf.kl_penalized_rewards(logp, ref, torch.tensor([1.5]), mask, beta=0.1)
    want = [-0.05, 0.02, 1.50, 0.0]
    assert got[0].tolist() == pytest.approx(want, abs=1e-6), (
        f"the chapter's table gives r_t = {want}; got {got[0].tolist()}")


# ---------------------------------------------------------------- Task 4: GAE over tokens
@pytest.mark.parametrize("lam,want_adv", [
    (1.0, [0.5, 0.4, 0.2]),    # lambda = 1: Monte Carlo return minus V
    (0.0, [0.1, 0.2, 0.2]),    # lambda = 0: one-step TD errors
    (0.5, [0.25, 0.3, 0.2]),
])
def test_task4_gae_hand_case(lam, want_adv):
    rewards = torch.tensor([[0.0, 0.0, 1.0]])
    values = torch.tensor([[0.5, 0.6, 0.8]])
    mask = torch.ones(1, 3)
    adv, ret = rlhf.gae_per_token(rewards, values, mask, gamma=1.0, lam=lam)
    assert adv[0].tolist() == pytest.approx(want_adv, abs=1e-6), f"lambda={lam}: expected advantages {want_adv}, got {adv[0].tolist()}"
    assert ret[0].tolist() == pytest.approx([a + v for a, v in zip(want_adv, [0.5, 0.6, 0.8])], abs=1e-6), "returns = advantages + values"


@pytest.mark.parametrize("lam,want_adv", [
    (1.0, [0.47, 0.32, 0.10]),       # Monte Carlo: return 1.47 minus V(s_1) = 1.0, etc.
    (0.95, [0.44925, 0.315, 0.10]),
])
def test_task4_chapter_paris_example(lam, want_adv):
    """The chapter's continued worked example: rewards (-0.05, 0.02, 1.50), values (1.0, 1.2, 1.4)."""
    rewards = torch.tensor([[-0.05, 0.02, 1.50]])
    values = torch.tensor([[1.0, 1.2, 1.4]])
    adv, _ = rlhf.gae_per_token(rewards, values, torch.ones(1, 3), gamma=1.0, lam=lam)
    assert adv[0].tolist() == pytest.approx(want_adv, abs=1e-5), (
        f"lambda={lam}: the chapter gives advantages {want_adv}, got {adv[0].tolist()}")


def test_task4_gae_ignores_padding():
    rewards = torch.tensor([[0.0, 1.0, 0.0, 0.0]])
    values = torch.tensor([[0.5, 0.6, 9.0, 9.0]])  # junk values at padding must be ignored
    mask = torch.tensor([[1.0, 1.0, 0.0, 0.0]])
    adv, ret = rlhf.gae_per_token(rewards, values, mask, gamma=1.0, lam=1.0)
    assert adv[0].tolist() == pytest.approx([0.5, 0.4, 0.0, 0.0], abs=1e-6), (
        f"got {adv[0].tolist()}: the state after the last real token has value 0, and padding gets advantage 0")
    assert ret[0].tolist() == pytest.approx([1.0, 1.0, 0.0, 0.0], abs=1e-6)


# ---------------------------------------------------------------- Task 5: PPO loss
def _loss_case():
    rho = torch.tensor([[1.1, 1.5, 0.5, 1.0]])
    old = torch.full((1, 4), math.log(0.5))
    logp = torch.log(0.5 * rho).requires_grad_(True)
    adv = torch.tensor([[2.0, 2.0, -1.0, 5.0]])
    values = torch.tensor([[1.0, 2.0, 3.0, 100.0]], requires_grad=True)
    returns = torch.tensor([[1.5, 2.0, 2.0, 0.0]])
    mask = torch.tensor([[1.0, 1.0, 1.0, 0.0]])
    return logp, old, values, adv, returns, mask


def test_task5_ppo_loss_value_and_stats():
    logp, old, values, adv, returns, mask = _loss_case()
    loss, stats = rlhf.ppo_rlhf_loss(logp, old, values, adv, returns, mask, clip_eps=0.2, vf_coef=0.5)
    # objective per real token: min(2.2, 2.2)=2.2, min(3.0, 2.4)=2.4, min(-0.5, -0.8)=-0.8 -> mean 1.2667
    # value loss: 0.5 * mean(0.25, 0, 1) = 0.20833
    assert stats["policy_loss"] == pytest.approx(-3.8 / 3, abs=1e-5), f"policy loss should be -1.26667, got {stats['policy_loss']}"
    assert stats["value_loss"] == pytest.approx(0.5 * 1.25 / 3, abs=1e-5), f"value loss should be 0.20833, got {stats['value_loss']}"
    assert loss.item() == pytest.approx(-3.8 / 3 + 0.5 * 0.5 * 1.25 / 3, abs=1e-5)
    assert stats["clip_frac"] == pytest.approx(2 / 3, abs=1e-6), "2 of the 3 real tokens have |rho-1| > 0.2"


def test_task5_ppo_loss_gradients():
    logp, old, values, adv, returns, mask = _loss_case()
    loss, _ = rlhf.ppo_rlhf_loss(logp, old, values, adv, returns, mask, clip_eps=0.2, vf_coef=0.5)
    loss.backward()
    g = logp.grad[0].tolist()
    assert g[0] == pytest.approx(-2.0 * 1.1 / 3, abs=1e-5), "inside the band: gradient -A*rho/N"
    assert g[1] == pytest.approx(0.0, abs=1e-7), "A > 0 and rho > 1+eps: clipped, gradient must be exactly 0"
    assert g[2] == pytest.approx(0.0, abs=1e-7), "A < 0 and rho < 1-eps: clipped, gradient must be exactly 0"
    assert g[3] == pytest.approx(0.0, abs=1e-7), "a padding token must get no gradient"
    assert values.grad[0, 3].item() == pytest.approx(0.0, abs=1e-7), "the value at a padding position must get no gradient"


# ---------------------------------------------------------------- end to end
def test_task5_rlhf_smoke_run_raises_reward():
    """A 12-iteration PPO run against a fixed 'even digits good' RM must raise the RM score."""
    sft = load_or_train_sft(os.path.join(HERE, "sft_tinylm.pt"))
    rm = DigitRM()
    with torch.no_grad():
        rm.w.copy_(torch.tensor([0.5, -0.5] * 5))
    cfg = rlhf.RLHFConfig(iters=12, beta=0.1, seed=0)
    _, hist = rlhf.train_rlhf(sft, rm, 0.0, cfg)
    start, end = sum(hist["rm"][:3]) / 3, sum(hist["rm"][-3:]) / 3
    assert end > start + 0.15, f"RM score should rise during PPO: first iterations {start:.2f}, last {end:.2f}"
    assert max(hist["malformed"][-3:]) < 0.2, "with beta = 0.1 and 12 iterations, answers should still be well formed"