"""Tests for Lab 16. Run one task at a time with  python -m pytest -q -k task1"""

import importlib
import math
import os

import pytest
import torch

from common import PRESETS, PGConfig, preset, train_length_toy

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

# The chapter's toy batch: two prompts, G = 4. Prompt 0 rewards 1,0,0,1; prompt 1 rewards 0,0,0,1.
R = torch.tensor([1.0, 0, 0, 1, 0, 0, 0, 1])
GID = torch.tensor([0, 0, 0, 0, 1, 1, 1, 1])


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
# ---------------------------------------------------------------------------
def test_task1_group_mean():
    A = zoo.group_advantages(R, GID, "group_mean", "none")
    want = [0.5, -0.5, -0.5, 0.5, -0.25, -0.25, -0.25, 0.75]
    assert close(A, want), f"group mean baseline: expected {want}, got {A.tolist()}"


def test_task1_leave_one_out():
    A = zoo.group_advantages(R, GID, "loo", "none")
    want = [2 / 3, -2 / 3, -2 / 3, 2 / 3, -1 / 3, -1 / 3, -1 / 3, 1.0]
    assert close(A, want), f"RLOO baseline (mean of the OTHER rewards): expected {want}, got {A.tolist()}"


def test_task1_group_std():
    A = zoo.group_advantages(R, GID, "group_mean", "group_std")
    s = math.sqrt(0.25 * 0.75)  # population std of (0,0,0,1)
    want = [1, -1, -1, 1, -0.25 / s, -0.25 / s, -0.25 / s, 0.75 / s]
    assert close(A, want, 1e-3), f"GRPO advantages (population std): expected {want}, got {A.tolist()}"


def test_task1_none_and_batch_mean():
    A = zoo.group_advantages(R, GID, "none", "none")
    assert close(A, R), "baseline 'none' should return the rewards unchanged"
    A = zoo.group_advantages(R, GID, "batch_mean", "none")
    assert close(A, R - 0.375), f"batch mean of these rewards is 0.375; got {A.tolist()}"


def test_task1_batch_std():
    A = zoo.group_advantages(R, GID, "group_mean", "batch_std")
    a = torch.tensor([0.5, -0.5, -0.5, 0.5, -0.25, -0.25, -0.25, 0.75])
    want = (a - a.mean()) / a.std(unbiased=False)
    assert close(A, want, 1e-3), f"REINFORCE++ batch normalisation: expected {want.tolist()}, got {A.tolist()}"


def test_task1_group_mean_is_rescaled_rloo():
    g = torch.Generator().manual_seed(0)
    r = torch.rand(12, generator=g)
    gid = torch.arange(3).repeat_interleave(4)
    a_mean = zoo.group_advantages(r, gid, "group_mean", "none")
    a_loo = zoo.group_advantages(r, gid, "loo", "none")
    assert close(a_mean, 0.75 * a_loo), "r_i - mean should equal (G-1)/G * (r_i - leave-one-out mean), with G = 4"


def test_task1_zero_variance_group_is_zero_and_finite():
    r = torch.tensor([1.0, 1, 1, 1, 0, 1, 0, 0])
    for base in ["group_mean", "loo"]:
        A = zoo.group_advantages(r, GID, base, "group_std")
        assert torch.isfinite(A).all(), "dividing by a zero std must stay finite (that is what eps is for)"
        assert close(A[:4], [0, 0, 0, 0]), f"an all-correct group carries no signal under {base}; got {A[:4].tolist()}"


# ---------------------------------------------------------------------------
# Task 2
# ---------------------------------------------------------------------------
X = torch.tensor([[1.0, 2.0, 0.0, 0.0], [1.0, 1.0, 1.0, 5.0]])
M = torch.tensor([[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0]])


def test_task2_values():
    assert close(zoo.aggregate(X, M, "seq_mean"), (1.5 + 2.0) / 2), "seq_mean: mean of per-response means (1.5, 2.0)"
    assert close(zoo.aggregate(X, M, "token_mean"), 11 / 6), "token_mean: (1+2+1+1+1+5) / 6 tokens"
    assert close(zoo.aggregate(X, M, "constant", max_len=4), 11 / 8), "constant: sum / (N * max_len) = 11 / 8"


def test_task2_ignores_padding():
    X2 = X.clone()
    X2[0, 2:] = 100.0  # garbage in padded positions
    for agg in ["seq_mean", "token_mean", "constant"]:
        assert close(zoo.aggregate(X2, M, agg, max_len=4), zoo.aggregate(X, M, agg, max_len=4)), \
            f"{agg}: values at mask == 0 must not count"


def test_task2_length_bias_in_gradients():
    x = torch.zeros(2, 4, requires_grad=True)
    zoo.aggregate(x, M, "seq_mean").backward()
    g = x.grad
    assert close(g[0, 0], 1 / (2 * 2)) and close(g[1, 0], 1 / (2 * 4)), \
        "seq_mean: each token of response i should get weight 1/(N |y_i|): 1/4 for the 2-token, 1/8 for the 4-token"
    for agg in ["token_mean", "constant"]:
        x = torch.zeros(2, 4, requires_grad=True)
        zoo.aggregate(x, M, agg, max_len=4).backward()
        assert close(x.grad[0, 0], x.grad[1, 0]), f"{agg}: every real token should get the same weight"


# ---------------------------------------------------------------------------
# Task 3
# ---------------------------------------------------------------------------
def grad_of_surrogate(rho, adv, cfg, mask=None):
    """Gradient of token_surrogate(...).sum() w.r.t. logp, at old_logp = log(0.5) and logp = old + log(rho)."""
    rho = torch.as_tensor(rho, dtype=torch.float32).reshape(1, -1)
    mask = torch.ones_like(rho) if mask is None else mask
    old = torch.full_like(rho, math.log(0.5))
    logp = (old + torch.log(rho)).clone().requires_grad_(True)
    term = zoo.token_surrogate(logp, old, torch.tensor([float(adv)]), mask, cfg)
    term.sum().backward()
    return term.detach(), logp.grad[0]


def test_task3_on_policy_all_equal():
    for name in ["grpo", "dapo", "cispo", "gspo", "reinforce"]:
        _, g = grad_of_surrogate([1.0, 1.0, 1.0], 0.7, PRESETS[name])
        assert close(g, [0.7, 0.7, 0.7]), f"{name}: on-policy (rho = 1) every token's gradient should be A = 0.7; got {g.tolist()}"


def test_task3_ppo_four_cases_with_clip_higher():
    cfg = PGConfig(clip="ppo", eps_low=0.2, eps_high=0.28)
    # A > 0: rho 1.1 inside -> grad rho*A; rho 1.4 above 1.28 -> 0; rho 0.5 below band -> still rho*A
    _, g = grad_of_surrogate([1.1, 1.4, 0.5], 1.0, cfg)
    assert close(g, [1.1, 0.0, 0.5]), f"A>0: expected gradients [1.1, 0, 0.5] (only rho > 1+eps_high is cut); got {g.tolist()}"
    # A < 0: rho 0.9 inside -> rho*A; rho 0.6 below 0.8 -> 0; rho 1.5 above band -> still rho*A
    _, g = grad_of_surrogate([0.9, 0.6, 1.5], -1.0, cfg)
    assert close(g, [-0.9, 0.0, -1.5]), f"A<0: expected [-0.9, 0, -1.5] (only rho < 1-eps_low is cut); got {g.tolist()}"


def test_task3_ppo_values():
    cfg = PGConfig(clip="ppo", eps_low=0.2, eps_high=0.2)
    term, _ = grad_of_surrogate([1.5, 0.5], 2.0, cfg)
    assert close(term[0], [2.4, 1.0]), f"min(rho*A, clip(rho)*A) with A=2: expected [2.4, 1.0]; got {term[0].tolist()}"


def test_task3_cispo_keeps_every_gradient():
    cfg = PRESETS["cispo"]  # eps_low = 1.0 (no lower bound), eps_high = 0.28
    _, g = grad_of_surrogate([1.5, 1.1, 0.3], 1.0, cfg)
    assert close(g, [1.28, 1.1, 0.3]), f"CISPO weight = sg(clip(rho, 0, 1.28)) * A, never zero; got {g.tolist()}"
    _, g = grad_of_surrogate([0.3], -1.0, cfg)
    assert close(g, [-0.3]), "CISPO with A<0 and rho = 0.3: weight 0.3 * (-1), not zero"


def test_task3_sequence_ratio_gspo():
    cfg = PGConfig(ratio="sequence", clip="none")
    rho = [1.3, 0.8, 1.0]
    s = (1.3 * 0.8 * 1.0) ** (1 / 3)
    term, g = grad_of_surrogate(rho, 1.0, cfg)
    assert close(term[0], [s, s, s]), f"sequence ratio = geometric mean {s:.4f} broadcast to every token; got {term[0].tolist()}"
    assert close(g, [s, s, s]), "d(sum of the 3 copies of s*A)/d logp_t = s*A for every token"
    # GSPO clip: the whole sequence is in or out together
    cfg = PGConfig(ratio="sequence", clip="ppo", eps_low=3e-4, eps_high=4e-4)
    _, g = grad_of_surrogate(rho, 1.0, cfg)
    assert close(g, [0, 0, 0]), f"s = {s:.4f} > 1 + 4e-4 with A > 0: GSPO clips every token of the response"


def test_task3_sequence_ratio_respects_mask():
    cfg = PGConfig(ratio="sequence", clip="none")
    mask = torch.tensor([[1.0, 1.0, 0.0]])
    term, _ = grad_of_surrogate([1.21, 1.0, 7.0], 1.0, cfg, mask)
    assert close(term[0, :2], [1.1, 1.1]), "geometric mean over the 2 real tokens only: sqrt(1.21) = 1.1"
    assert close(term[0, 2], 0.0), "padded positions must be zero"


# ---------------------------------------------------------------------------
# Task 4
# ---------------------------------------------------------------------------
def test_task4_estimators():
    logp = torch.log(torch.tensor([0.5, 0.2]))
    ref = torch.log(torch.tensor([0.25, 0.4]))
    d = ref - logp
    assert close(zoo.kl_per_token(logp, ref, "k1"), -d), "k1 = log(pi/pi_ref)"
    assert close(zoo.kl_per_token(logp, ref, "k2"), 0.5 * d ** 2), "k2 = (log(pi/pi_ref))^2 / 2"
    want = torch.tensor([0.5 - math.log(0.5) - 1, 2 - math.log(2) - 1])
    assert close(zoo.kl_per_token(logp, ref, "k3"), want), f"k3 = pi_ref/pi - log(pi_ref/pi) - 1: expected {want.tolist()}"
    assert (zoo.kl_per_token(logp, ref, "k3") >= 0).all(), "k3 is never negative"


# ---------------------------------------------------------------------------
# Task 5: each preset against a hand-written formula from its paper
# ---------------------------------------------------------------------------
def fixed_batch(seed=0):
    g = torch.Generator().manual_seed(seed)
    N, T = 8, 6
    lens = torch.tensor([2, 6, 3, 5, 1, 4, 6, 2])
    mask = (torch.arange(T)[None, :] < lens[:, None]).float()
    old = -torch.rand(N, T, generator=g) * 2 - 0.1
    logp = old + 0.3 * torch.randn(N, T, generator=g)  # a few updates in: ratios spread around 1
    ref = old + 0.2 * torch.randn(N, T, generator=g)
    return logp, old, ref, mask


def reference_objective(name, logp, old, ref, rewards, mask, gid):
    """Each algorithm written the way its paper writes it, with explicit loops."""
    cfg = PRESETS[name]
    N = len(rewards)
    lens = mask.sum(1)
    groups = sorted(set(gid.tolist()))
    members = {g: [i for i in range(N) if gid[i] == g] for g in groups}

    def mean_std(idx):
        v = torch.stack([rewards[j] for j in idx])
        return v.mean(), v.std(unbiased=False)

    total = 0.0
    if name in ("reinforce", "rloo"):
        for i in range(N):
            if name == "reinforce":
                A = rewards[i]
            else:
                others = [rewards[j] for j in members[int(gid[i])] if j != i]
                A = rewards[i] - sum(others) / len(others)
            for t in range(int(lens[i])):
                rho = torch.exp(logp[i, t] - old[i, t])
                total = total + rho * A
        return total / N
    if name == "grpo":  # DeepSeekMath eq. 3: (1/G) sum_i (1/|y_i|) sum_t {min(...) - beta k3}
        for i in range(N):
            m, s = mean_std(members[int(gid[i])])
            A = (rewards[i] - m) / (s + 1e-6)
            inner = 0.0
            for t in range(int(lens[i])):
                rho = torch.exp(logp[i, t] - old[i, t])
                d = ref[i, t] - logp[i, t]
                k3 = torch.exp(d) - d - 1
                inner = inner + torch.min(rho * A, torch.clamp(rho, 0.8, 1.2) * A) - cfg.beta * k3
            total = total + inner / lens[i]
        return total / N
    if name == "dr_grpo":  # no std, no 1/|y_i|: constant normaliser
        for i in range(N):
            m, _ = mean_std(members[int(gid[i])])
            A = rewards[i] - m
            for t in range(int(lens[i])):
                rho = torch.exp(logp[i, t] - old[i, t])
                total = total + torch.min(rho * A, torch.clamp(rho, 0.8, 1.2) * A)
        return total / (N * cfg.max_len)
    if name in ("dapo", "cispo"):  # 1 / sum_i |y_i| outside both sums
        for i in range(N):
            m, s = mean_std(members[int(gid[i])])
            A = (rewards[i] - m) / (s + 1e-6)
            for t in range(int(lens[i])):
                rho = torch.exp(logp[i, t] - old[i, t])
                if name == "dapo":
                    total = total + torch.min(rho * A, torch.clamp(rho, 0.8, 1.28) * A)
                else:
                    total = total + torch.clamp(rho, 0.0, 1.28).detach() * A * logp[i, t]
        return total / lens.sum()
    if name == "gspo":  # (1/G) sum_i min(s_i A_i, clip(s_i) A_i)
        for i in range(N):
            m, s = mean_std(members[int(gid[i])])
            A = (rewards[i] - m) / (s + 1e-6)
            si = torch.exp(sum(logp[i, t] - old[i, t] for t in range(int(lens[i]))) / lens[i])
            total = total + torch.min(si * A, torch.clamp(si, 1 - 3e-4, 1 + 4e-4) * A)
        return total / N
    raise KeyError(name)


@pytest.mark.parametrize("name", ["reinforce", "rloo", "grpo", "dr_grpo", "dapo", "gspo", "cispo"])
def test_task5_presets_match_their_papers(name):
    logp, old, ref, mask = fixed_batch()
    rewards = torch.tensor([1.0, 0, 0, 1, 0, 0, 0, 1])
    lp1 = logp.clone().requires_grad_(True)
    loss = zoo.pg_loss(lp1, old, ref, rewards, mask, GID, PRESETS[name])
    loss.backward()
    lp2 = logp.clone().requires_grad_(True)
    want = -reference_objective(name, lp2, old, ref, rewards, mask, GID)
    want.backward()
    assert close(loss.detach(), want.detach(), 1e-4), f"{name}: loss {loss.item():.6f}, expected {want.item():.6f}"
    assert torch.allclose(lp1.grad, lp2.grad, atol=1e-5), f"{name}: gradient w.r.t. logp differs from the paper's formula"


def test_task5_zero_variance_group_gives_zero_gradient():
    logp, old, ref, mask = fixed_batch(1)
    rewards = torch.tensor([1.0, 1, 1, 1, 0, 1, 0, 0])  # group 0 all correct
    for name in ["rloo", "dr_grpo", "dapo", "gspo", "cispo"]:
        lp = logp.clone().requires_grad_(True)
        zoo.pg_loss(lp, old, None, rewards, mask, GID, PRESETS[name]).backward()
        assert lp.grad[:4].abs().max() < 1e-6, f"{name}: an all-correct group must not move the policy"
        assert lp.grad[4:].abs().max() > 1e-4, f"{name}: the mixed group should still get a gradient"


def test_task5_kl_placements():
    logp, old, ref, mask = fixed_batch(2)
    rewards = torch.tensor([1.0, 0, 0, 1, 0, 0, 0, 1])
    base = PGConfig(baseline="none", scale="none", aggregation="token_mean", clip="none")
    no_kl = zoo.pg_loss(logp, old, ref, rewards, mask, GID, base)
    cfg = preset("grpo", baseline="none", scale="none", aggregation="token_mean", clip="none",
                 kl="loss", kl_estimator="k2", beta=0.5)
    in_loss = zoo.pg_loss(logp, old, ref, rewards, mask, GID, cfg)
    k2 = (0.5 * (ref - logp) ** 2 * mask).sum() / mask.sum()
    assert close(in_loss - no_kl, 0.5 * k2), "KL in the loss: loss should grow by beta * token-mean of k2"
    cfg = preset("grpo", baseline="none", scale="none", aggregation="constant", max_len=1, clip="none",
                 kl="reward", beta=0.1)
    lp = old.clone().requires_grad_(True)  # on-policy
    zoo.pg_loss(lp, old, ref, rewards, mask, GID, cfg).backward()
    shaped = rewards - 0.1 * ((old - ref) * mask).sum(1)
    assert torch.allclose(lp.grad, -(shaped[:, None] * mask) / 8, atol=1e-5), \
        "KL in the reward: each token's gradient should be the KL-shaped sequence reward (no extra gradient path)"


# ---------------------------------------------------------------------------
# Does it learn? And does it show the length bias? (LengthToy, a few seconds)
# ---------------------------------------------------------------------------
def test_task5_learns_on_length_toy():
    h = train_length_toy(zoo.pg_loss, PRESETS["dr_grpo"], steps=150)
    start, end = sum(h["reward"][:10]) / 10, sum(h["reward"][-10:]) / 10
    assert end > 0.85 and end > start + 0.3, f"Dr. GRPO should solve the toy: mean reward {start:.2f} -> {end:.2f}"


def test_task5_seq_mean_makes_wrong_answers_longer():
    def drift(name):
        h = train_length_toy(zoo.pg_loss, preset(name, kl="none"), steps=120)
        w = [x for x in h["len_wrong"] if x == x]
        return sum(w[-20:]) / 20 - sum(w[:20]) / 20

    d_grpo, d_dr = drift("grpo"), drift("dr_grpo")
    assert d_grpo > 2.0, f"GRPO's seq-mean should make wrong answers longer (drift {d_grpo:+.2f} tokens)"
    assert abs(d_dr) < 1.5, f"Dr. GRPO should not (drift {d_dr:+.2f} tokens)"