"""Tests for policy.py. Run: python -m pytest -q   (or -k task1, -k task3, ...)"""

import importlib
import os
import random

import numpy as np
import pytest
import torch
import torch.nn.functional as F

import tinylm
import toylm

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


# =============================================================================== Task 1
def test_task1_shift_on_a_hand_example():
    # vocabulary of 3, one sequence of length 4: prompt = ids[:2], response = ids[2:]
    ids = torch.tensor([[0, 1, 2, 0]])
    logits = torch.zeros(1, 4, 3)
    logits[0, 1] = torch.tensor([0.0, 0.0, 2.0])   # column 1 predicts ids[2] = 2
    logits[0, 2] = torch.tensor([1.0, 0.0, 0.0])   # column 2 predicts ids[3] = 0
    logits[0, 3] = torch.tensor([9.0, -9.0, 5.0])  # column 3 predicts a token that doesn't exist: ignored
    out = policy.response_logprobs(logits, ids, prompt_len=2)
    assert out.shape == (1, 2), f"expected one log-prob per response token, shape (1, 2), got {tuple(out.shape)}"
    want = torch.tensor([[2 - np.log(2 + np.exp(2)), 1 - np.log(np.e + 2)]], dtype=torch.float32)
    assert torch.allclose(out, want, atol=1e-5), (
        f"expected {want.tolist()}, got {out.tolist()}. "
        "Remember: logits at column t give the distribution of the token at column t + 1.")


def test_task1_prompt_logits_do_not_matter():
    torch.manual_seed(0)
    ids = torch.tensor([[3, 1, 4, 1, 5]])
    logits = torch.randn(1, 5, 7)
    a = policy.response_logprobs(logits, ids, prompt_len=3)
    logits2 = logits.clone()
    logits2[0, 0] += 100 * torch.randn(7)  # column 0 predicts a prompt token
    logits2[0, 1] += 100 * torch.randn(7)  # column 1 predicts a prompt token
    b = policy.response_logprobs(logits2, ids, prompt_len=3)
    assert torch.allclose(a, b), "changing predictions of PROMPT tokens must not change the response log-probs"


def test_task1_matches_tinylm_on_real_samples():
    torch.manual_seed(1)
    model = tinylm.TinyLM()
    prompts, _ = tinylm.make_batch(8, random.Random(0))
    out = tinylm.sample(model, prompts)
    with torch.no_grad():
        logits = model(out["ids"])
        mine = policy.response_logprobs(logits, out["ids"], out["prompt_len"])
        ref = tinylm.sequence_logprobs(model, out)
    m = out["mask"].bool()
    assert mine.shape == ref.shape, f"shape {tuple(mine.shape)}, expected {tuple(ref.shape)}"
    assert torch.allclose(mine[m], ref[m], atol=1e-5), "per-token log-probs disagree with tinylm.sequence_logprobs on response tokens"


def test_task1b_sequence_logprob_sums_and_masks():
    tok = torch.tensor([[-0.5, -0.25, -1.0], [-2.0, float("-inf"), float("nan")]])
    mask = torch.tensor([[1.0, 1.0, 1.0], [1.0, 0.0, 0.0]])
    s = policy.sequence_logprob(tok, mask)
    assert s.shape == (2,), f"expected shape (2,), got {tuple(s.shape)}"
    assert torch.isfinite(s).all(), "masked positions leaked into the sum (0 * -inf = nan): use torch.where"
    assert torch.allclose(s, torch.tensor([-1.75, -2.0])), f"expected [-1.75, -2.0], got {s.tolist()}"


def test_task1b_gradient_reaches_only_response_tokens():
    torch.manual_seed(2)
    model = tinylm.TinyLM()
    prompts, _ = tinylm.make_batch(4, random.Random(1))
    out = tinylm.sample(model, prompts)
    logits = model(out["ids"])
    logits.retain_grad()
    tok = policy.response_logprobs(logits, out["ids"], out["prompt_len"])
    policy.sequence_logprob(tok, out["mask"]).sum().backward()
    L = out["prompt_len"]
    g = logits.grad.abs().sum(-1)  # (B, L_total)
    assert g[:, : L - 1].sum() == 0, "prompt tokens received gradient: only response tokens should"
    assert g[:, L - 1:].sum() > 0, "response tokens received no gradient"


# =============================================================================== Task 2
def test_task2_greedy_at_zero_temperature():
    logits = torch.tensor([[0.1, 2.0, -1.0], [5.0, 4.9, float("-inf")]])
    out = policy.sample_with_temperature(logits, 0.0)
    assert out.tolist() == [1, 0], f"T = 0 must be argmax, got {out.tolist()}"
    assert out.dtype == torch.long


@pytest.mark.parametrize("T", [0.5, 1.0, 2.0])
def test_task2_frequencies_match_softmax(T):
    z = torch.tensor([2.0, 1.0, 0.0, float("-inf")])
    g = torch.Generator().manual_seed(0)
    n = 20000
    draws = policy.sample_with_temperature(z.expand(n, 4), T, g)
    assert draws.shape == (n,), f"expected shape ({n},), got {tuple(draws.shape)}"
    freq = torch.bincount(draws, minlength=4).float() / n
    want = F.softmax(z / T, dim=-1)
    assert freq[3] == 0, "a token with logit -inf must never be sampled"
    assert torch.allclose(freq, want, atol=0.015), f"T={T}: frequencies {freq.tolist()} vs softmax(z/T) {want.tolist()}"


def test_task2_reproducible_with_generator():
    z = torch.randn(5, 10)
    a = policy.sample_with_temperature(z, 1.0, torch.Generator().manual_seed(3))
    b = policy.sample_with_temperature(z, 1.0, torch.Generator().manual_seed(3))
    assert torch.equal(a, b), "pass the generator to torch.multinomial so sampling is reproducible"


def test_task2_generate_runs_the_token_mdp():
    torch.manual_seed(0)
    model = tinylm.TinyLM()
    prompts, _ = tinylm.make_batch(6, random.Random(0))
    out = policy.generate(model, prompts, temperature=1.0, generator=torch.Generator().manual_seed(0))
    assert out["ids"].shape == (6, out["prompt_len"] + tinylm.MAX_ANSWER)
    greedy = policy.generate(model, prompts, temperature=0.0)
    ref = tinylm.sample(model, prompts, greedy=True)
    assert torch.equal(greedy["ids"], ref["ids"]), "greedy generation should match tinylm.sample(greedy=True)"


# =============================================================================== Task 3
P3 = torch.tensor([0.6, 0.3, 0.1], dtype=torch.float64)
Q3 = torch.tensor([0.4, 0.4, 0.2], dtype=torch.float64)
EXACT_KL = float((P3 * (P3 / Q3).log()).sum())  # 0.08766


def test_task3_values_on_the_worked_example():
    k1, k2, k3 = policy.kl_estimators(P3.log(), Q3.log())
    assert torch.allclose(k1, torch.tensor([0.4055, -0.2877, -0.6931], dtype=torch.float64), atol=1e-4), f"k1 = {k1.tolist()}"
    assert torch.allclose(k2, torch.tensor([0.0822, 0.0414, 0.2402], dtype=torch.float64), atol=1e-4), f"k2 = {k2.tolist()}"
    assert torch.allclose(k3, torch.tensor([0.0721, 0.0457, 0.3069], dtype=torch.float64), atol=1e-4), f"k3 = {k3.tolist()}"


def test_task3_exact_expectations():
    # Expectation under pi (= P3), computed exactly by weighting each token by its probability.
    k1, k2, k3 = policy.kl_estimators(P3.log(), Q3.log())
    e1, e2, e3 = [float((P3 * k).sum()) for k in (k1, k2, k3)]
    assert e1 == pytest.approx(EXACT_KL, abs=1e-9), "k1 is unbiased: E_pi[k1] = KL"
    assert e3 == pytest.approx(EXACT_KL, abs=1e-9), "k3 is unbiased: E_pi[k3] = KL, because E_pi[pi_ref/pi] = 1"
    assert abs(e2 - EXACT_KL) > 1e-3, "k2 should be (slightly) biased"


def test_task3_k3_never_negative_and_k1_sometimes_is():
    g = torch.Generator().manual_seed(0)
    a, b = torch.randn(10000, generator=g), torch.randn(10000, generator=g)
    k1, k2, k3 = policy.kl_estimators(a, b)
    assert (k3 >= -1e-7).all(), "every single k3 value must be >= 0 (log u <= u - 1)"
    assert (k2 >= 0).all()
    assert (k1 < 0).any(), "k1 can be negative for a single sample"


def test_task3_monte_carlo_on_sampled_tokens():
    g = torch.Generator().manual_seed(1)
    y = torch.multinomial(P3.float(), 200000, replacement=True, generator=g)
    k1, k2, k3 = policy.kl_estimators(P3.log()[y], Q3.log()[y])
    assert float(k1.mean()) == pytest.approx(EXACT_KL, abs=0.004)
    assert float(k3.mean()) == pytest.approx(EXACT_KL, abs=0.002)
    assert float(k3.std()) < 0.25 * float(k1.std()), "k3 should have far lower variance than k1 here"


# =============================================================================== Task 4
W = toylm.SCENARIOS["worked"]
W_REF = toylm.seq_probs(W["first"], W["second"])


def test_task4_worked_example_from_the_chapter():
    pi, log_Z = policy.tilted_optimum(W_REF, W["reward"], 1.0)
    want = np.array([[0.1722, 0.3120], [0.0615, 0.4543]])
    assert np.allclose(pi, want, atol=1e-4), f"expected {want.round(4).tolist()}, got {np.round(pi, 4).tolist()}"
    assert log_Z == pytest.approx(np.log(2.43947), abs=1e-4)


def test_task4_limits_and_stability():
    big, _ = policy.tilted_optimum(W_REF, W["reward"], 1e4)
    assert np.allclose(big, W_REF, atol=1e-3), "beta -> infinity: pi* -> pi_ref"
    small, log_Z = policy.tilted_optimum(W_REF, W["reward"], 1e-3)
    assert np.all(np.isfinite(small)) and np.isfinite(log_Z), "overflow: subtract the max before exponentiating"
    assert small[1, 1] == pytest.approx(1.0, abs=1e-6), "beta -> 0: all mass on the best response"


def test_task4_zero_reference_mass_stays_zero():
    ref = np.array([[0.0, 0.5], [0.25, 0.25]])
    r = np.array([[100.0, 0.0], [0.0, 1.0]])
    pi, _ = policy.tilted_optimum(ref, r, 0.5)
    assert pi[0, 0] == 0.0, "a response the reference never produces stays impossible, whatever its reward"


def test_task4_objective_values():
    pi, log_Z = policy.tilted_optimum(W_REF, W["reward"], 1.0)
    assert policy.kl_objective(pi, W_REF, W["reward"], 1.0) == pytest.approx(log_Z, abs=1e-9), "J(pi*) = beta log Z"
    assert policy.kl_objective(W_REF, W_REF, W["reward"], 1.0) == pytest.approx(0.58), "J(pi_ref) = E_ref[r] = 0.58"
    greedy = np.array([[0, 0], [0, 1.0]])
    assert policy.kl_objective(greedy, W_REF, W["reward"], 1.0) == pytest.approx(2 - np.log(1 / 0.15)), "0 log 0 = 0"


def test_task4_nothing_beats_pi_star():
    rng = np.random.default_rng(0)
    for _ in range(5):
        first, second, r = toylm.random_toy_lm(3, rng)
        ref = toylm.seq_probs(first, second)
        beta = float(rng.uniform(0.1, 3))
        pi, log_Z = policy.tilted_optimum(ref, r, beta)
        best = policy.kl_objective(pi, ref, r, beta)
        for _ in range(300):
            other = rng.dirichlet(np.ones(9)).reshape(3, 3)
            mix = 0.5 * other + 0.5 * pi
            for q in (other, mix):
                gap = best - policy.kl_objective(q, ref, r, beta)
                # The Gibbs identity: J(q) = beta log Z - beta KL(q || pi*)
                assert gap == pytest.approx(beta * toylm.kl(q, pi), abs=1e-8), "J(pi*) - J(q) should equal beta * KL(q || pi*)"
                assert gap >= -1e-12


# =============================================================================== Task 5
def test_task5_worked_example_soft_values():
    f, s, V, V0 = policy.token_level_optimum(W["first"], W["second"], W["reward"], 1.0)
    assert np.allclose(V, [0.5231, 1.4338], atol=1e-4), f"soft values V(a), V(b): expected [0.5231, 1.4338], got {np.round(V, 4)}"
    assert V0 == pytest.approx(0.8918, abs=1e-4), "V0 = beta log Z"
    assert np.allclose(f, [0.4842, 0.5158], atol=1e-4), f"first-token pi*: expected [0.4842, 0.5158], got {np.round(f, 4)}"
    assert np.allclose(s, [[0.3556, 0.6444], [0.1192, 0.8808]], atol=1e-4)


@pytest.mark.parametrize("seed", [0, 1, 2])
def test_task5_agrees_with_sequence_level(seed):
    rng = np.random.default_rng(seed)
    first, second, r = toylm.random_toy_lm(4, rng)
    for beta in (0.05, 0.7, 5.0):
        f, s, V, V0 = policy.token_level_optimum(first, second, r, beta)
        pi, log_Z = policy.tilted_optimum(toylm.seq_probs(first, second), r, beta)
        assert np.allclose(f[:, None] * s, pi, atol=1e-9), "pi*(i) pi*(j|i) must equal the sequence-level pi*(ij)"
        assert V0 == pytest.approx(beta * log_Z, rel=1e-9, abs=1e-9)
        assert np.allclose(s.sum(1), 1) and f.sum() == pytest.approx(1)