"""Tests for Lab 15. Run one task at a time with:  python -m pytest -q -k task2"""

import copy
import importlib
import math
import os
import random

import pytest
import torch

import tinylm
from tinylm import TinyLM, reward, sample, sequence_logprobs

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

LAB = os.path.dirname(os.path.abspath(__file__))
SFT_PATH = os.path.join(LAB, "tinylm_sft.pt")


def close(a, b, tol=1e-4):
    return torch.allclose(torch.as_tensor(a, dtype=torch.float32), torch.as_tensor(b, dtype=torch.float32), atol=tol)


# ---------------------------------------------------------------- Task 1: sample_group
def test_task1_shapes_and_grouping():
    torch.manual_seed(0)
    model = TinyLM()
    prompts, answers = ["1+2=", "30+40=", "99+99="], ["3", "70", "198"]
    G = 4
    g = grpo.sample_group(model, prompts, answers, G)
    out = g["out"]
    assert out["ids"].shape[0] == 12, f"expected P*G = 12 rows, got {out['ids'].shape[0]}"
    assert len(g["responses"]) == 12, "expected one decoded response per row"
    assert g["rewards"].shape == (12,), f"rewards should have shape (12,), got {tuple(g['rewards'].shape)}"
    assert g["rewards"].dtype == torch.float32, "rewards should be a float32 tensor"
    L = out["prompt_len"]
    for i, p in enumerate(prompts):
        block = out["ids"][i * G:(i + 1) * G, :L]
        assert (block == block[0]).all(), f"rows {i*G}..{(i+1)*G-1} should all hold prompt {p!r} (repeat each prompt G times, consecutively)"
        assert tinylm.decode(block[0].tolist()) == p, f"rows {i*G}..{(i+1)*G-1} should hold prompt {p!r}"


def test_task1_rewards_match_verifier():
    torch.manual_seed(1)
    model = TinyLM()
    rng = random.Random(0)
    prompts, answers = tinylm.make_batch(5, rng)
    G = 3
    g = grpo.sample_group(model, prompts, answers, G)
    rep_answers = [a for a in answers for _ in range(G)]
    expect = [reward(r, a) for r, a in zip(g["responses"], rep_answers)]
    assert g["rewards"].tolist() == expect, "rewards[i] should be tinylm.reward(responses[i], answer of row i)"
    assert g["responses"] == tinylm.decode_responses(g["out"]), "responses should be decode_responses(out)"


# ---------------------------------------------------------------- Task 2: group_advantages
def test_task2_chapter_example():
    # The chapter's worked example: rewards 1, 0, 0, 1  ->  advantages +1, -1, -1, +1
    adv = grpo.group_advantages(torch.tensor([1.0, 0.0, 0.0, 1.0]), G=4)
    assert close(adv, [1, -1, -1, 1], tol=1e-3), f"rewards (1,0,0,1): expected about (1,-1,-1,1), got {adv.tolist()}"


def test_task2_one_right_three_wrong():
    adv = grpo.group_advantages(torch.tensor([1.0, 0.0, 0.0, 0.0]), G=4)
    s = math.sqrt(3)
    assert close(adv, [s, -1 / s, -1 / s, -1 / s], tol=1e-3), (
        f"rewards (1,0,0,0): mean 0.25, population std 0.433, so expected "
        f"(1.732, -0.577, -0.577, -0.577); got {adv.tolist()}. Did you use the population std (unbiased=False)?")


def test_task2_groups_are_independent_and_all_equal_gives_zero():
    r = torch.tensor([1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0])
    adv = grpo.group_advantages(r, G=4)
    assert adv.shape == (12,), f"expected shape (12,), got {tuple(adv.shape)}"
    assert torch.isfinite(adv).all(), "all-equal groups must not produce NaN or inf (that is what eps is for)"
    assert close(adv[:8], torch.zeros(8)), f"all-right and all-wrong groups should get advantage 0, got {adv[:8].tolist()}"
    assert close(adv[8:], [1, -1, -1, 1], tol=1e-3), "each group must be normalized with its own mean and std"


def test_task2_mean_only():
    adv = grpo.group_advantages(torch.tensor([1.0, 0.0, 0.0, 0.0, 2.0, 4.0]), G=3, scale_std=False)
    expect = [2 / 3, -1 / 3, -1 / 3, -2.0, 0.0, 2.0]
    assert close(adv, expect), f"scale_std=False: expected r - group mean = {expect}, got {adv.tolist()}"


# ---------------------------------------------------------------- Task 3: k3_kl
def test_task3_hand_value_and_gradient():
    logp = torch.tensor([[math.log(0.5)]], requires_grad=True)
    ref = torch.tensor([[math.log(0.25)]])
    k = grpo.k3_kl(logp, ref)
    expect = 0.5 - math.log(0.5) - 1  # 0.19315
    assert close(k, [[expect]]), f"pi = 0.5, pi_ref = 0.25: expected k3 = {expect:.5f}, got {k.item():.5f}"
    k.sum().backward()
    assert logp.grad is not None, "the gradient must flow through logp"
    assert close(logp.grad, [[0.5]]), f"d k3 / d logp = 1 - pi_ref/pi = 0.5; got {logp.grad.item():.4f}"


def test_task3_nonnegative_and_zero_at_equality():
    g = torch.Generator().manual_seed(0)
    logp = torch.randn(50, 7, generator=g) - 2
    ref = torch.randn(50, 7, generator=g) - 2
    k = grpo.k3_kl(logp, ref)
    assert k.shape == (50, 7), "k3 must be computed per token (same shape as the inputs)"
    assert (k >= -1e-7).all(), "k3 is never negative (x - 1 - log x >= 0)"
    assert close(grpo.k3_kl(logp, logp), torch.zeros(50, 7)), "k3 must be 0 when the two models agree"


def test_task3_unbiased_for_kl():
    pi = torch.tensor([0.6, 0.3, 0.1])
    ref = torch.tensor([0.3, 0.3, 0.4])
    exact = float((pi * (pi / ref).log()).sum())
    g = torch.Generator().manual_seed(0)
    a = torch.multinomial(pi, 200_000, replacement=True, generator=g)
    est = grpo.k3_kl(pi.log()[a], ref.log()[a]).mean().item()
    assert abs(est - exact) < 0.01, f"averaged over samples from pi, k3 should estimate KL(pi||ref) = {exact:.4f}; got {est:.4f}"


# ---------------------------------------------------------------- Task 4: grpo_loss
def _hand_batch():
    # Two responses. Response 0: A = +1, two real tokens with ratios 1.5 and 0.5.
    # Response 1: A = -1, three real tokens with ratios 1.5, 0.5, 1.0.
    old = torch.zeros(2, 3)
    ratios = torch.tensor([[1.5, 0.5, 1.0], [1.5, 0.5, 1.0]])
    logp = ratios.log().clone().requires_grad_(True)
    mask = torch.tensor([[1.0, 1.0, 0.0], [1.0, 1.0, 1.0]])
    adv = torch.tensor([1.0, -1.0])
    return logp, old, adv, mask


def test_task4_hand_value_and_clip_fraction():
    logp, old, adv, mask = _hand_batch()
    loss, stats = grpo.grpo_loss(logp, old, logp.detach(), adv, mask, clip_eps=0.2, beta=0.0)
    # response 0: min(1.5, 1.2) = 1.2 and min(0.5, 0.8) = 0.5 -> mean 0.85
    # response 1: min(-1.5, -1.2) = -1.5, min(-0.5, -0.8) = -0.8, -1.0 -> mean -1.1
    assert close(loss, 0.125), f"expected loss = -((0.85 + -1.1) / 2) = 0.125, got {float(loss):.4f}"
    assert abs(stats["clip_frac"] - 0.4) < 1e-6, f"2 of the 5 real tokens are clipped; expected clip_frac 0.4, got {stats['clip_frac']}"
    assert abs(stats["ratio_max"] - 1.5) < 1e-5, f"expected ratio_max 1.5, got {stats['ratio_max']}"


def test_task4_gradients_vanish_exactly_where_clipped():
    logp, old, adv, mask = _hand_batch()
    loss, _ = grpo.grpo_loss(logp, old, logp.detach(), adv, mask, clip_eps=0.2, beta=0.0)
    loss.backward()
    # d loss / d logp = -(rho * A) / (|y_i| * B) where not clipped, 0 where clipped, 0 on padding
    expect = torch.tensor([[0.0, -0.125, 0.0], [0.25, 0.0, 1 / 6]])
    assert close(logp.grad, expect), (
        f"expected d loss / d logp =\n{expect}\ngot\n{logp.grad}\n"
        "(clipped tokens and padding get zero gradient; each response is averaged over its own length)")


def test_task4_padding_is_ignored():
    logp, old, adv, mask = _hand_batch()
    junk = logp.detach().clone()
    junk[0, 2] = 5.0  # garbage at a padded position
    l1, _ = grpo.grpo_loss(logp.detach(), old, logp.detach(), adv, mask, beta=0.04)
    l2, _ = grpo.grpo_loss(junk, old, junk, adv, mask, beta=0.04)
    assert close(l1, l2), "values at padded positions (mask = 0) must not change the loss"


def test_task4_kl_term():
    # ratio 1, advantage 0: only the KL term is left.  k3(pi=0.5, ref=0.25) = 0.19315
    logp = torch.full((1, 2), math.log(0.5))
    ref = torch.full((1, 2), math.log(0.25))
    loss, stats = grpo.grpo_loss(logp, logp, ref, torch.zeros(1), torch.ones(1, 2), beta=0.1)
    k = 0.5 - math.log(0.5) - 1
    assert close(loss, 0.1 * k), f"expected loss = beta * k3 = {0.1 * k:.5f}, got {float(loss):.5f} (note the sign: we minimize -objective)"
    assert abs(stats["kl"] - k) < 1e-5, f"stats['kl'] should be the mean k3 over real tokens, {k:.5f}"


def test_task4_ratio_one_is_reinforce_with_group_baseline():
    torch.manual_seed(0)
    logp = (torch.randn(4, 5) - 1).requires_grad_(True)
    mask = torch.tensor([[1, 1, 1, 0, 0], [1, 1, 1, 1, 1], [1, 0, 0, 0, 0], [1, 1, 0, 0, 0]], dtype=torch.float32)
    adv = torch.tensor([0.5, -1.0, 2.0, -1.5])
    loss, _ = grpo.grpo_loss(logp, logp.detach(), logp.detach(), adv, mask, beta=0.0)
    loss.backward()
    expect = -(adv[:, None] / mask.sum(1, keepdim=True)) * mask / 4
    assert close(logp.grad, expect), "at ratio 1 the gradient must be -A_i / (|y_i| * B) on every real token"


def test_task4_zero_signal_group_gives_zero_gradient():
    torch.manual_seed(0)
    model = TinyLM()
    out = sample(model, ["12+30="] * 4 + ["5+5="] * 4)
    rewards = torch.tensor([1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 1.0])
    adv = grpo.group_advantages(rewards, G=4)
    logp = sequence_logprobs(model, out)
    zero = slice(0, 4)  # the all-correct group only
    loss, _ = grpo.grpo_loss(logp[zero], logp[zero].detach(), logp[zero].detach(), adv[zero], out["mask"][zero], beta=0.0)
    loss.backward()
    total = sum(float(p.grad.abs().sum()) for p in model.parameters() if p.grad is not None)
    assert total < 1e-8, f"a group whose rewards are all equal must give exactly zero gradient; got total |grad| = {total:.3g}"


# ---------------------------------------------------------------- Task 5: grpo_step and training
def _small_cfg(**kw):
    cfg = grpo.GRPOConfig(prompts_per_iter=4, G=4, inner_steps=2, lr=1e-3)
    for k, v in kw.items():
        setattr(cfg, k, v)
    return cfg


def test_task5_step_updates_policy_not_reference():
    torch.manual_seed(0)
    policy = TinyLM()
    ref = TinyLM()  # a different network, so the KL term has a non-zero gradient
    for p in ref.parameters():
        p.requires_grad_(False)
    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)
    prompts, answers = ["1+1=", "2+2=", "3+4=", "50+50="], ["2", "4", "7", "100"]
    s = grpo.grpo_step(policy, ref, opt, prompts, answers, _small_cfg(beta=0.1))
    for key in ["reward", "loss", "kl", "clip_frac", "frac_zero_signal", "entropy"]:
        assert key in s, f"grpo_step should return a dict with key {key!r}"
        assert isinstance(s[key], float), f"stats[{key!r}] should be a Python float"
    assert all(torch.equal(a, b) for a, b in zip(ref_before, ref.parameters())), "the reference model must stay frozen"
    assert any(not torch.equal(a, b) for a, b in zip(pol_before, policy.parameters())), "the policy should change (beta > 0 and policy != reference, so the KL term has a gradient)"


def test_task5_one_inner_step_never_clips():
    torch.manual_seed(0)
    policy = TinyLM()
    ref = copy.deepcopy(policy)
    opt = torch.optim.Adam(policy.parameters(), lr=1e-3)
    prompts, answers = ["1+1=", "2+2=", "3+4=", "50+50="], ["2", "4", "7", "100"]
    s = grpo.grpo_step(policy, ref, opt, prompts, answers, _small_cfg(inner_steps=1))
    assert s["clip_frac"] == 0.0, (
        "with one inner step the ratio is exactly 1 (old_logp comes from the same weights), so nothing is clipped. "
        "Did you compute old_logp before the update, with the policy itself?")


def test_task5_grpo_improves_the_sft_model():
    cfg = grpo.GRPOConfig(iters=100, eval_every=100, eval_n=500, lr=3e-4, seed=0)
    policy0 = tinylm.load_or_train_sft(SFT_PATH)
    before = tinylm.accuracy(policy0, n=500, greedy=False)
    policy, hist = grpo.train_grpo(cfg, sft_path=SFT_PATH, verbose=False)
    after = hist["acc_sampled"][-1]
    assert after >= before + 0.06, (
        f"100 GRPO iterations should raise sampled accuracy by at least 6 points (the reference solution gains 12-18); "
        f"went {before:.3f} -> {after:.3f}")
    assert len(hist["kl"]) == 100 and all(k >= 0 for k in hist["kl"]), "KL (k3) should be logged every iteration and be >= 0"