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

import copy
import importlib
import math
import os
import random

import pytest
import torch

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

from prefs import response_logprobs  # noqa: E402
from tinylm import TinyLM, make_batch  # noqa: E402


def t(x):
    return torch.tensor(x, dtype=torch.float32)


# ---------------------------------------------------------------- Task 1
def test_task1_sum_ignores_padding():
    tok = t([[-1.0, -2.0, -0.5, 99.0], [-0.3, -0.7, 5.0, 5.0]])
    mask = t([[1, 1, 1, 0], [1, 1, 0, 0]])
    out = dpo.sequence_logprob(tok, mask)
    assert out.shape == (2,), f"expected shape (2,), got {tuple(out.shape)}"
    assert torch.allclose(out, t([-3.5, -1.0])), (
        f"expected [-3.5, -1.0], got {out.tolist()}. Padding positions (mask 0) hold junk like 99: "
        "multiply by the mask before summing.")


def test_task1_average_divides_by_length():
    tok = t([[-1.0, -2.0, -0.5, 99.0], [-0.3, -0.7, 5.0, 5.0]])
    mask = t([[1, 1, 1, 0], [1, 1, 0, 0]])
    out = dpo.sequence_logprob(tok, mask, average=True)
    assert torch.allclose(out, t([-3.5 / 3, -0.5])), (
        f"expected [-3.5/3, -0.5] (sum divided by the number of real tokens), got {out.tolist()}")


def test_task1_gradient_flows():
    tok = t([[-1.0, -2.0, 0.0]]).requires_grad_(True)
    mask = t([[1, 1, 0]])
    dpo.sequence_logprob(tok, mask).sum().backward()
    assert torch.allclose(tok.grad, mask), "d(sum)/d(token_logp) should equal the mask"


# ---------------------------------------------------------------- Task 2
# The chapter's worked example: beta = 0.1,
#   log pi_theta(y_w) = -12, log pi_ref(y_w) = -13, log pi_theta(y_l) = -15, log pi_ref(y_l) = -14
#   r_w = 0.1, r_l = -0.1, margin 0.2, loss = -log sigmoid(0.2) = 0.598139, weight sigmoid(-0.2) = 0.450166
def test_task2_implicit_reward():
    r = dpo.implicit_reward(t([-12.0, -15.0]), t([-13.0, -14.0]), 0.1)
    assert torch.allclose(r, t([0.1, -0.1])), f"expected [0.1, -0.1], got {r.tolist()}"


def test_task2_dpo_loss_worked_example():
    pw = t([-12.0]).requires_grad_(True)
    pl = t([-15.0]).requires_grad_(True)
    loss, r_w, r_l = dpo.dpo_loss(pw, pl, t([-13.0]), t([-14.0]), 0.1)
    assert loss.dim() == 0, "the loss must be a scalar (use .mean())"
    assert loss.item() == pytest.approx(0.598139, abs=1e-5), (
        f"expected -log sigmoid(0.2) = 0.598139, got {loss.item():.6f}")
    assert r_w.item() == pytest.approx(0.1, abs=1e-6) and r_l.item() == pytest.approx(-0.1, abs=1e-6)
    assert not r_w.requires_grad and not r_l.requires_grad, "return the implicit rewards detached"
    loss.backward()
    w = 0.1 * (1 / (1 + math.exp(0.2)))  # beta * sigmoid(-margin)
    assert pw.grad.item() == pytest.approx(-w, abs=1e-6), (
        f"d loss / d log pi(y_w) should be -beta*sigmoid(r_l - r_w) = {-w:.6f}, got {pw.grad.item():.6f}")
    assert pl.grad.item() == pytest.approx(+w, abs=1e-6), (
        f"d loss / d log pi(y_l) should be +beta*sigmoid(r_l - r_w) = {w:.6f}, got {pl.grad.item():.6f}")


def test_task2_dpo_loss_is_log2_at_reference():
    lp = t([-3.0, -7.0, -1.5])
    loss, r_w, r_l = dpo.dpo_loss(lp, lp - 1.0, lp, lp - 1.0, 0.5)
    assert loss.item() == pytest.approx(math.log(2), abs=1e-6), (
        f"when the policy equals the reference every margin is 0 and the loss is log 2 = 0.6931; got {loss.item():.6f}")


def test_task2_dpo_loss_batch_mean_and_stability():
    pw = t([0.0, -500.0]).requires_grad_(True)
    loss, _, _ = dpo.dpo_loss(pw, t([0.0, 0.0]), t([0.0, 0.0]), t([0.0, 0.0]), 1.0)
    # margins 0 and -500: losses log 2 and 500 (approximately); mean 250.3466
    assert math.isfinite(loss.item()), "the loss overflowed: use F.logsigmoid"
    assert loss.item() == pytest.approx((math.log(2) + 500.0) / 2, abs=1e-3), (
        f"expected the batch MEAN of [log 2, 500], got {loss.item()}")
    loss.backward()
    assert torch.isfinite(pw.grad).all(), "NaN/inf gradient: use F.logsigmoid"


# ---------------------------------------------------------------- Task 3
def test_task3_ipo_loss_value_and_gradient():
    pw = t([-12.0]).requires_grad_(True)
    loss = dpo.ipo_loss(pw, t([-15.0]), t([-13.0]), t([-14.0]), 0.5)
    # h = (-12+13) - (-15+14) = 2, target 1/(2*0.5) = 1, loss (2-1)^2 = 1
    assert loss.item() == pytest.approx(1.0, abs=1e-6), f"expected (h - 1/(2 beta))^2 = (2 - 1)^2 = 1, got {loss.item()}"
    loss.backward()
    assert pw.grad.item() == pytest.approx(2.0, abs=1e-6), f"expected d loss / d log pi(y_w) = 2(h - 1) = 2, got {pw.grad.item()}"


def test_task3_ipo_minimum_at_finite_margin():
    beta = 0.2
    target = 1 / (2 * beta)  # 2.5
    loss = dpo.ipo_loss(t([target]), t([0.0]), t([0.0]), t([0.0]), beta)
    assert loss.item() == pytest.approx(0.0, abs=1e-8), "IPO's loss is zero when the log-ratio margin equals 1/(2 beta)"
    bigger = dpo.ipo_loss(t([10 * target]), t([0.0]), t([0.0]), t([0.0]), beta)
    assert bigger.item() > 1.0, "IPO should PENALIZE margins beyond the target (that's the point of IPO)"


def test_task3_simpo_loss():
    loss = dpo.simpo_loss(t([-0.5, -0.5]), t([-1.5, -1.5]), 2.0, 0.5)
    # z = 2*(-0.5) - 2*(-1.5) - 0.5 = 1.5; loss = log(1 + e^-1.5) = 0.201413
    assert loss.dim() == 0
    assert loss.item() == pytest.approx(0.201413, abs=1e-5), f"expected 0.201413, got {loss.item():.6f}"
    zero_margin = dpo.simpo_loss(t([-0.5]), t([-1.5]), 2.0, 0.0)
    assert zero_margin.item() < loss.item(), "a positive margin should make the loss LARGER (a harder target)"


# ---------------------------------------------------------------- Task 4
def _pairs(n, seed=0, max_n=20):
    """Chosen = the correct sum, rejected = a wrong sum (off by 1..3)."""
    rng = random.Random(seed)
    prompts, answers = make_batch(n, rng, max_n)
    return [{"prompt": p, "chosen": a, "rejected": str(int(a) + rng.randint(1, 3))}
            for p, a in zip(prompts, answers)]


def _models(seed=0):
    torch.manual_seed(seed)
    policy = TinyLM()
    ref = copy.deepcopy(policy)
    for p in ref.parameters():
        p.requires_grad_(False)
    return policy, ref


def test_task4_first_step_matches_hand_values():
    policy, ref = _models()
    batch = _pairs(8)
    opt = torch.optim.Adam(policy.parameters(), lr=1e-3)
    stats = dpo.dpo_step(policy, ref, batch, opt, dpo.DPOConfig(loss="dpo", beta=0.5))
    for k in ["loss", "reward_acc", "margin", "chosen_logp", "rejected_logp"]:
        assert k in stats and isinstance(stats[k], float), f"stats must contain float '{k}'"
    assert stats["loss"] == pytest.approx(math.log(2), abs=1e-5), (
        "on the first step the policy equals the reference, so the DPO loss must be log 2 = 0.6931; "
        f"got {stats['loss']:.6f}. Did you compute the reference log-probs with the ref model?")
    assert stats["margin"] == pytest.approx(0.0, abs=1e-5)
    assert stats["reward_acc"] == 0.0, "with all margins exactly 0, no pair has r_w > r_l"
    with torch.no_grad():
        tw, mw = response_logprobs(ref, [b["prompt"] for b in batch], [b["chosen"] for b in batch])
        expected_chosen = (tw * mw).sum(-1).mean().item()
    assert stats["chosen_logp"] == pytest.approx(expected_chosen, abs=1e-4), (
        "chosen_logp should be the batch mean of log pi_theta(y_w|x) BEFORE the update")


def test_task4_reference_is_frozen_and_policy_moves():
    policy, ref = _models()
    ref_before = [p.clone() for p in ref.parameters()]
    pol_before = [p.clone() for p in policy.parameters()]
    opt = torch.optim.Adam(policy.parameters(), lr=1e-3)
    dpo.dpo_step(policy, ref, _pairs(8), opt, dpo.DPOConfig(loss="dpo", beta=0.5))
    assert all(torch.equal(a, b) for a, b in zip(ref_before, ref.parameters())), "the reference model must not change"
    assert all(p.grad is None for p in ref.parameters()), "no gradient should reach the reference (use torch.no_grad())"
    assert any(not torch.equal(a, b) for a, b in zip(pol_before, policy.parameters())), "the policy should have been updated"


def test_task4_nll_ipo_simpo_first_step_values():
    batch = _pairs(8)
    prompts = [b["prompt"] for b in batch]
    with torch.no_grad():
        policy, ref = _models()
        tw, mw = response_logprobs(policy, prompts, [b["chosen"] for b in batch])
        tl, ml = response_logprobs(policy, prompts, [b["rejected"] for b in batch])
        avg_w = (tw * mw).sum(-1) / mw.sum(-1)
        avg_l = (tl * ml).sum(-1) / ml.sum(-1)

    # DPO + NLL: log 2 + c * mean per-token NLL of the chosen
    policy, ref = _models()
    opt = torch.optim.Adam(policy.parameters(), lr=1e-3)
    s = dpo.dpo_step(policy, ref, batch, opt, dpo.DPOConfig(loss="dpo", beta=0.5, nll_coef=0.2))
    exp = math.log(2) - 0.2 * avg_w.mean().item()
    assert s["loss"] == pytest.approx(exp, abs=1e-4), f"DPO + 0.2*NLL at step 0: expected {exp:.5f}, got {s['loss']:.5f}"

    # IPO: margin 0, loss (0 - 1/(2 beta))^2 = 1 for beta = 0.5
    policy, ref = _models()
    opt = torch.optim.Adam(policy.parameters(), lr=1e-3)
    s = dpo.dpo_step(policy, ref, batch, opt, dpo.DPOConfig(loss="ipo", beta=0.5))
    assert s["loss"] == pytest.approx(1.0, abs=1e-5), f"IPO at step 0 with beta=0.5: expected 1.0, got {s['loss']:.5f}"

    # SimPO: no reference; uses length-normalized policy log-probs
    policy, ref = _models()
    opt = torch.optim.Adam(policy.parameters(), lr=1e-3)
    s = dpo.dpo_step(policy, ref, batch, opt, dpo.DPOConfig(loss="simpo", beta=2.0, simpo_margin=0.5))
    z = 2.0 * avg_w - 2.0 * avg_l - 0.5
    exp = torch.nn.functional.softplus(-z).mean().item()
    assert s["loss"] == pytest.approx(exp, abs=1e-4), f"SimPO at step 0: expected {exp:.5f}, got {s['loss']:.5f}"

    with pytest.raises(ValueError):
        policy, ref = _models()
        opt = torch.optim.Adam(policy.parameters(), lr=1e-3)
        dpo.dpo_step(policy, ref, batch, opt, dpo.DPOConfig(loss="ppo"))


def test_task4_learns_the_preferences():
    """A short DPO run on 32 fixed pairs: the implicit reward should rank every pair correctly."""
    policy, ref = _models(seed=1)
    pairs = _pairs(32, seed=3)
    opt = torch.optim.Adam(policy.parameters(), lr=1e-3)
    cfg = dpo.DPOConfig(loss="dpo", beta=0.5)
    first = None
    for _ in range(40):
        s = dpo.dpo_step(policy, ref, pairs, opt, cfg)
        first = first if first is not None else s
    assert s["loss"] < 0.3, f"after 40 steps on the same 32 pairs the DPO loss should be < 0.3 (started at 0.693); got {s['loss']:.3f}"
    assert s["reward_acc"] >= 0.95, f"reward accuracy should approach 1, got {s['reward_acc']:.2f}"
    assert s["margin"] > first["margin"] + 1.0, "the implicit-reward margin should grow"