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

import importlib
import os

import numpy as np
import pytest

SOL = os.environ.get("RL_SOLUTIONS")
memory = importlib.import_module("solutions.memory" if SOL else "memory")
scheduler = importlib.import_module("solutions.scheduler" if SOL else "scheduler")
tis = importlib.import_module("solutions.tis" if SOL else "tis")

from scheduler_core import lognormal_lengths  # noqa: E402
from toy import estimator_moments, make_toy, sample_from_sampler  # noqa: E402


# ======================================================================== Task 1
def test_task1a_model_state_bytes():
    assert memory.model_state_bytes(7e9, trainable=True) == pytest.approx(112e9), (
        "a trained 7B model: 16 bytes/param (bf16 weights 2 + bf16 grads 2 + fp32 master 4 + Adam m 4 + v 4)")
    assert memory.model_state_bytes(7e9, trainable=False) == pytest.approx(14e9), (
        "a frozen 7B model in bf16: 2 bytes/param")
    assert memory.model_state_bytes(1e9, trainable=False, bytes_per_weight=1) == pytest.approx(1e9), (
        "bytes_per_weight should be used for frozen models (e.g. 1 byte for FP8)")


def test_task1b_kv_cache_bytes():
    per_token = memory.kv_cache_bytes(n_layers=28, n_kv_heads=4, head_dim=128, n_tokens=1)
    assert per_token == 57344, (
        f"Qwen2.5-7B shape: 2 (K and V) * 28 * 4 * 128 * 2 bytes = 57,344 bytes per token, got {per_token}")
    big = memory.kv_cache_bytes(28, 4, 128, n_tokens=64 * 16384)
    assert big == pytest.approx(64 * 16384 * 57344), "KV cache is linear in the number of tokens"
    assert memory.kv_cache_bytes(28, 4, 128, 1, bytes_per_value=1) == 28672, "FP8 KV cache halves it"


def test_task1c_ppo_vs_grpo():
    ppo = memory.rl_memory_gb("ppo", 7e9)
    for key in ("policy", "reference", "critic", "reward_model", "total"):
        assert key in ppo, f"missing key {key!r} in the returned dict"
    expected = {"policy": 112, "reference": 14, "critic": 112, "reward_model": 14, "total": 252}
    for k, v in expected.items():
        assert ppo[k] == pytest.approx(v), f"PPO 7B: expected {k} = {v} GB, got {ppo[k]}"

    grpo = memory.rl_memory_gb("grpo", 7e9)
    assert grpo["total"] == pytest.approx(126), (
        f"GRPO 7B with a reference model and a verifier: 112 + 14 = 126 GB, got {grpo['total']}")
    assert grpo["critic"] == 0 and grpo["reward_model"] == 0, "GRPO has no critic, and a verifier costs no GPU memory"
    assert memory.rl_memory_gb("grpo", 7e9, use_ref=False)["total"] == pytest.approx(112), "beta = 0: no reference"
    assert memory.rl_memory_gb("grpo", 7e9, rm_params=7e9)["total"] == pytest.approx(140), (
        "GRPO with a learned 7B reward model: 112 + 14 + 14")
    small = memory.rl_memory_gb("ppo", 7e9, critic_params=1e9, rm_params=1e9)
    assert small["critic"] == pytest.approx(16) and small["reward_model"] == pytest.approx(2)


def test_task1c_lora():
    lora = memory.rl_memory_gb("grpo", 7e9, lora_params=40e6)
    assert lora["policy"] == pytest.approx(14 + 0.64), (
        f"LoRA: frozen base 7e9 * 2 B = 14 GB + trainable adapter 40e6 * 16 B = 0.64 GB; got {lora['policy']}")
    assert lora["reference"] == 0, "with LoRA the reference model is the base model (adapter off): no extra copy"
    assert lora["total"] == pytest.approx(14.64)


def test_task1c_bad_algo():
    with pytest.raises(ValueError):
        memory.rl_memory_gb("dqn", 7e9)


# ======================================================================== Task 2
def test_task2a_may_start():
    ms = scheduler.may_start
    assert ms(0, 0, 0, 0, 4) and ms(3, 0, 0, 0, 4), "sync: at version 0 you may start responses 0..3"
    assert not ms(4, 0, 0, 0, 4), "sync (k=0): the 5th response must wait for the trainer"
    assert ms(4, 0, 1, 0, 4), "after one training step (version 1), the next batch may start"
    assert ms(7, 0, 0, 1, 4) and not ms(8, 0, 0, 1, 4), "k=1: up to 2 batches may be started at version 0"
    assert ms(8, 1, 0, 1, 4), "a dropped response frees its slot"


def _q(*tags):
    return [{"id": i, "tag": t} for i, t in enumerate(tags)]


def test_task2b_take_batch_fifo_and_lag():
    q = _q(0, 1, 1, 2)
    batch, dropped, rest = scheduler.take_batch(q, version=2, batch_size=2, max_staleness=5, drop_stale=False)
    assert batch is not None and [s["id"] for s in batch] == [0, 1], "take the FIRST B samples, in queue order"
    assert [s["lag"] for s in batch] == [2, 1], "lag = version - tag"
    assert dropped == [] and [s["id"] for s in rest] == [2, 3]


def test_task2b_take_batch_not_enough():
    q = _q(0)
    batch, dropped, rest = scheduler.take_batch(q, version=0, batch_size=2, max_staleness=0, drop_stale=False)
    assert batch is None, "fewer than B samples: return None for the batch"
    assert [s["id"] for s in rest] == [0], "and keep the queue intact"


def test_task2b_take_batch_drop():
    q = _q(0, 0, 1)
    batch, dropped, rest = scheduler.take_batch(q, version=2, batch_size=2, max_staleness=1, drop_stale=True)
    assert batch is None, "after dropping the two lag-2 samples only one is left"
    assert sorted(s["id"] for s in dropped) == [0, 1], "samples with lag > k are dropped"
    assert [s["id"] for s in rest] == [2]


def test_task2_sync_hand_example():
    # 2 lanes, batch 2, train_time 1. Step 0: lengths 1 and 3 -> batch done at t=3, train until 4.
    # Step 1: lengths 1 and 1 -> done at 5, train until 6.  Lanes busy 6 of 2*6 = 12 time units.
    r = scheduler.simulate([1, 3, 1, 1] + [1] * 10, n_lanes=2, batch_size=2, max_staleness=0,
                           train_time=1.0, n_steps=2)
    assert r["wall_time"] == pytest.approx(6.0), f"sync: expected wall time 6, got {r['wall_time']}"
    assert r["lane_util"] == pytest.approx(0.5), f"expected lane utilization 0.5, got {r['lane_util']}"
    assert r["mean_lag"] == 0 and r["max_lag"] == 0, "synchronous RL is exactly on-policy"


def test_task2_async_hand_example():
    # Same lengths, k = 1: the short responses of the next batch fill the gaps.
    r = scheduler.simulate([1, 3, 1, 1] + [1] * 10, n_lanes=2, batch_size=2, max_staleness=1,
                           train_time=1.0, n_steps=2)
    assert r["wall_time"] == pytest.approx(4.0), f"async k=1: expected wall time 4, got {r['wall_time']}"
    assert r["lane_util"] == pytest.approx(1.0)
    assert r["mean_lag"] == pytest.approx(0.5), "step 1 trains on two samples tagged version 0: lags [0, 0, 1, 1]"


def test_task2_long_tail_story():
    L = lognormal_lengths(20000, sigma=1.0, seed=0)
    sync = scheduler.simulate(L, 32, 32, 0, 1.0, 40)
    asyn = scheduler.simulate(L, 32, 32, 2, 1.0, 40)
    assert sync["lane_util"] < 0.3, f"sync with a heavy tail should waste most lane time, util = {sync['lane_util']:.2f}"
    assert asyn["throughput"] > 3 * sync["throughput"], "async (k=2) should be several times faster here"
    assert asyn["max_lag"] > 2, "admission control bounds how far ahead generation runs, not the lag of long stragglers"
    dropped = scheduler.simulate(L, 32, 32, 2, 1.0, 40, drop_stale=True)
    assert dropped["max_lag"] <= 2 and dropped["n_dropped"] > 0


# ======================================================================== Task 3
def test_task3a_weights_hand_values():
    lt = np.log(np.array([[0.5, 0.2]]))
    ls = np.log(np.array([[0.25, 0.4]]))      # per-token ratios 2.0 and 0.5; sequence ratio 1.0
    iw = tis.importance_weights
    np.testing.assert_allclose(iw(lt, ls, "none"), [[1, 1]], err_msg="'none': all ones")
    np.testing.assert_allclose(iw(lt, ls, "token"), [[2.0, 0.5]], err_msg="'token': pi(y_t)/mu(y_t)")
    np.testing.assert_allclose(iw(lt, ls, "token_tis", C=1.5), [[1.5, 0.5]], err_msg="'token_tis': min(w_t, C)")
    np.testing.assert_allclose(iw(lt, ls, "seq"), [[1.0, 1.0]], err_msg="'seq': prod_t w_t, same for every token")
    lt2 = np.log(np.array([[0.5, 0.4]]))      # ratios 2 and 1 -> sequence ratio 2
    np.testing.assert_allclose(iw(lt2, ls, "seq_tis", C=1.5), [[1.5, 1.5]], err_msg="'seq_tis': min(W, C)")
    with pytest.raises(ValueError):
        iw(lt, ls, "bogus")


def test_task3a_weights_shape_and_identity():
    toy = make_toy(mismatch=0.0)
    for mode in tis.MODES:
        w = tis.importance_weights(toy["logp_train"], toy["logp_sampler"], mode)
        assert w.shape == toy["logp_train"].shape, f"{mode}: expected shape (N, T)"
        np.testing.assert_allclose(w, 1.0, atol=1e-12, err_msg=f"{mode}: with no mismatch every weight is 1")


def test_task3b_pg_estimate_hand():
    rewards = np.array([1.0, 0.0])
    score = np.ones((2, 2, 3))
    weights = np.array([[2.0, 0.5], [1.0, 1.0]])
    est = tis.pg_estimate(rewards, score, weights, baseline=0.25)
    assert est.shape == (2, 2, 3), "per-sample estimates have shape (N, T, V)"
    np.testing.assert_allclose(est[0, 0], 0.75 * 2.0, err_msg="(r - b) * w_t * score")
    np.testing.assert_allclose(est[0, 1], 0.75 * 0.5)
    np.testing.assert_allclose(est[1], -0.25)


def _moments(toy, mode, C=2.0):
    w = tis.importance_weights(toy["logp_train"], toy["logp_sampler"], mode, C)
    return estimator_moments(toy, tis.pg_estimate(toy["R"], toy["score"], w))


def test_task3_on_policy_is_unbiased():
    toy = make_toy(mismatch=0.0)
    m = _moments(toy, "none")
    assert m["rel_bias"] < 1e-10, "with sampler == trainer, plain REINFORCE is unbiased"


def test_task3_mismatch_biases_naive_and_seq_is_fixes_it():
    toy = make_toy(mismatch=0.3)
    naive = _moments(toy, "none")
    seq = _moments(toy, "seq")
    tok = _moments(toy, "token_tis")
    assert naive["rel_bias"] > 0.2, f"naive estimator should be visibly biased, rel bias {naive['rel_bias']:.3f}"
    assert seq["rel_bias"] < 1e-8, "sequence-level importance sampling is exactly unbiased"
    assert tok["rel_bias"] < 0.5 * naive["rel_bias"], "per-token TIS should remove most of the bias"
    assert seq["var"] > tok["var"], "the price of exact IS is variance"


def test_task3_truncation_wins_at_large_mismatch():
    toy = make_toy(mismatch=0.6)
    rmse = {m: _moments(toy, m)["rel_rmse"](256) for m in tis.MODES}
    assert rmse["seq"] > 2 * rmse["seq_tis"], f"untruncated sequence IS should blow up: {rmse}"
    assert min(rmse["token_tis"], rmse["seq_tis"]) < rmse["none"], f"a truncated estimator should beat naive: {rmse}"


def test_task3_sampled_estimate_matches_exact_mean():
    toy = make_toy(mismatch=0.3)
    idx = sample_from_sampler(toy, 200000, np.random.default_rng(0))
    w = tis.importance_weights(toy["logp_train"][idx], toy["logp_sampler"][idx], "token_tis")
    est = tis.pg_estimate(toy["R"][idx], toy["score"][idx], w).mean(0)
    exact = _moments(toy, "token_tis")["mean"]
    assert np.abs(est - exact).max() < 0.02, "a big sampled batch should agree with the exact expectation"