"""Tests for opd.py (Tasks 1-3). Run: python -m pytest -q test_opd.py"""

import importlib
import math
import os

import torch
import torch.nn.functional as F

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

torch.set_num_threads(1)


# ---------------------------------------------------------------- Task 1
def test_reverse_kl_rewards_hand_case():
    # Two responses, two positions; second position of response 1 is padding.
    s = torch.log(torch.tensor([[0.5, 0.25], [0.1, 0.8]]))
    t = torch.log(torch.tensor([[0.25, 0.9], [0.4, 0.8]]))
    mask = torch.tensor([[1.0, 0.0], [1.0, 1.0]])
    r = opd.reverse_kl_rewards(s, t, mask)
    expected = torch.tensor([[math.log(0.25 / 0.5), 0.0], [math.log(0.4 / 0.1), 0.0]])
    assert torch.allclose(r, expected, atol=1e-6), f"expected log(pi_T/pi_theta) masked, got {r}"


def test_reverse_kl_rewards_is_detached():
    s = torch.log(torch.tensor([[0.5, 0.5]])).requires_grad_(True)
    t = torch.log(torch.tensor([[0.25, 0.75]]))
    r = opd.reverse_kl_rewards(s, t, torch.ones(1, 2))
    assert not r.requires_grad, "rewards must be detached (a reward is a constant, not part of the loss graph)"


# ---------------------------------------------------------------- Task 2
def test_opd_pg_loss_value_and_gradient():
    logp = torch.log(torch.tensor([[0.5, 0.25, 0.9], [0.2, 0.6, 0.3]])).requires_grad_(True)
    rewards = torch.tensor([[1.0, -2.0, 0.5], [3.0, 0.0, 7.0]])
    mask = torch.tensor([[1.0, 1.0, 0.0], [1.0, 1.0, 1.0]])
    loss = opd.opd_pg_loss(logp, rewards, mask)
    n = mask.sum()
    expected = -(rewards * logp.detach() * mask).sum() / n
    assert torch.allclose(loss, expected, atol=1e-6), f"loss {loss.item():.5f}, expected {expected.item():.5f}"
    loss.backward()
    assert torch.allclose(logp.grad, -rewards * mask / n, atol=1e-6), \
        "d loss / d logp[b,t] should be -r[b,t] * mask[b,t] / sum(mask)"


def test_opd_pg_loss_does_not_backprop_into_rewards():
    logp = torch.log(torch.tensor([[0.5, 0.5]])).requires_grad_(True)
    rewards = torch.tensor([[1.0, 2.0]], requires_grad=True)
    opd.opd_pg_loss(logp, rewards, torch.ones(1, 2)).backward()
    assert rewards.grad is None or torch.all(rewards.grad == 0), "rewards must be treated as constants (stop-gradient)"


# ---------------------------------------------------------------- Task 3
def _manual_kl(ls, lt):
    ps, pt = F.softmax(ls, -1), F.softmax(lt, -1)
    return (ps * (ps.log() - pt.log())).sum(-1)


def test_full_vocab_kl_value():
    torch.manual_seed(0)
    ls, lt = torch.randn(2, 3, 5), torch.randn(2, 3, 5)
    mask = torch.tensor([[1.0, 1.0, 0.0], [1.0, 0.0, 0.0]])
    got = opd.full_vocab_reverse_kl(ls, lt, mask)
    want = (_manual_kl(ls, lt) * mask).sum() / mask.sum()
    assert torch.allclose(got, want, atol=1e-5), f"got {got.item():.5f}, expected {want.item():.5f}"
    assert got.item() >= 0


def test_full_vocab_kl_zero_and_flat_when_equal():
    torch.manual_seed(1)
    ls = torch.randn(3, 4, 6).requires_grad_(True)
    lt = ls.detach().clone() + 5.0  # adding a constant to logits does not change the distribution
    loss = opd.full_vocab_reverse_kl(ls, lt, torch.ones(3, 4))
    assert abs(loss.item()) < 1e-6, "KL of a distribution with itself must be 0"
    loss.backward()
    assert ls.grad.abs().max() < 1e-6, "at the minimum (student == teacher) the gradient must vanish"


def test_full_vocab_kl_teacher_is_constant():
    ls = torch.randn(1, 2, 4).requires_grad_(True)
    lt = torch.randn(1, 2, 4).requires_grad_(True)
    opd.full_vocab_reverse_kl(ls, lt, torch.ones(1, 2)).backward()
    assert lt.grad is None or torch.all(lt.grad == 0), "the teacher is frozen: no gradient may flow into teacher_logits"


def test_sampled_token_reward_is_unbiased_for_full_kl():
    """E_{y ~ pi_theta}[-r] equals the exact reverse KL (Monte Carlo, fixed seed)."""
    torch.manual_seed(2)
    ls, lt = torch.randn(1, 1, 6), torch.randn(1, 1, 6)
    exact = opd.full_vocab_reverse_kl(ls, lt, torch.ones(1, 1)).item()
    ps = F.softmax(ls[0, 0], -1)
    y = torch.multinomial(ps, 200_000, replacement=True)
    s_logp = F.log_softmax(ls[0, 0], -1)[y][None]
    t_logp = F.log_softmax(lt[0, 0], -1)[y][None]
    r = opd.reverse_kl_rewards(s_logp, t_logp, torch.ones_like(s_logp))
    est = -r.mean().item()
    assert abs(est - exact) < 0.02, f"mean of -r over samples = {est:.4f}, exact KL = {exact:.4f}"


# ---------------------------------------------------------------- does it learn?
def test_sampled_token_opd_learns_a_teacher_distribution():
    """A one-token 'language model' (just a logits vector) trained by sampled-token OPD
    on its own samples should move most of the way to the teacher."""
    torch.manual_seed(3)
    teacher = torch.log(torch.tensor([0.05, 0.6, 0.25, 0.05, 0.05]))
    logits = torch.zeros(5, requires_grad=True)
    opt = torch.optim.Adam([logits], lr=0.05)
    kl0 = opd.full_vocab_reverse_kl(logits.detach()[None, None], teacher[None, None], torch.ones(1, 1)).item()
    for _ in range(400):
        y = torch.multinomial(F.softmax(logits, -1).detach(), 64, replacement=True)
        s_logp = F.log_softmax(logits, -1)[y][None]
        r = opd.reverse_kl_rewards(s_logp, teacher[y][None], torch.ones(1, 64))
        loss = opd.opd_pg_loss(s_logp, r, torch.ones(1, 64))
        opt.zero_grad()
        loss.backward()
        opt.step()
    kl1 = opd.full_vocab_reverse_kl(logits.detach()[None, None], teacher[None, None], torch.ones(1, 1)).item()
    assert kl1 < 0.1 * kl0, f"reverse KL should fall by 10x: started {kl0:.3f}, ended {kl1:.3f}"