"""Tests for tools.py (Tasks 4-5). Run: python -m pytest -q test_tools.py"""

import importlib
import math
import os
import random

import numpy as np
import pytest

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


# ---------------------------------------------------------------- Task 4
def test_pass_at_k_hand_values():
    assert math.isclose(tools.pass_at_k(4, 2, 1), 0.5), "n=4, c=2: pass@1 = c/n = 0.5"
    assert math.isclose(tools.pass_at_k(4, 2, 2), 5 / 6), "n=4, c=2: pass@2 = 1 - C(2,2)/C(4,2) = 5/6"
    assert tools.pass_at_k(10, 0, 5) == 0.0, "no correct samples: pass@k = 0"
    assert tools.pass_at_k(10, 7, 4) == 1.0, "n - c = 3 < k = 4: every 4-subset contains a correct sample"
    assert math.isclose(tools.pass_at_k(16, 3, 1), 3 / 16), "pass@1 is the fraction correct"


def test_pass_at_k_matches_brute_force():
    rng = random.Random(0)
    n, c, k = 8, 3, 3
    samples = [1] * c + [0] * (n - c)
    hits = sum(any(rng.sample(samples, k)) for _ in range(100_000)) / 100_000
    assert abs(tools.pass_at_k(n, c, k) - hits) < 0.01, "should equal P(at least one correct among k drawn without replacement)"


def test_pass_at_k_is_stable_for_large_n():
    v = tools.pass_at_k(1000, 3, 500)
    exact = 1 - math.comb(997, 500) / math.comb(1000, 500)
    assert math.isclose(v, exact, rel_tol=1e-9), f"got {v}, expected {exact}"
    assert isinstance(v, float)


def test_pass_at_k_rejects_bad_arguments():
    with pytest.raises(ValueError):
        tools.pass_at_k(4, 5, 1)
    with pytest.raises(ValueError):
        tools.pass_at_k(4, 2, 5)


# ---------------------------------------------------------------- Task 5
def test_maxrl_hand_values():
    np.testing.assert_allclose(tools.maxrl_advantages([1, 0, 0, 1]), [1, -1, -1, 1],
                               err_msg="p_hat = 0.5: (r - 0.5) / 0.5")
    np.testing.assert_allclose(tools.maxrl_advantages([1, 0, 0, 0]), [3, -1, -1, -1],
                               err_msg="p_hat = 0.25: (r - 0.25) / 0.25")


def test_maxrl_zero_signal_groups():
    np.testing.assert_allclose(tools.maxrl_advantages([0, 0, 0, 0]), [0, 0, 0, 0], err_msg="no success: zeros, no NaN")
    np.testing.assert_allclose(tools.maxrl_advantages([1, 1, 1]), [0, 0, 0], err_msg="all correct: zeros")


def test_maxrl_batched_groups_and_equivalence():
    r = np.array([[1, 0, 0, 0], [1, 1, 1, 0], [0, 0, 0, 0]], dtype=float)
    a = tools.maxrl_advantages(r)
    assert a.shape == r.shape
    np.testing.assert_allclose(a[1], [1 / 3, 1 / 3, 1 / 3, -1], err_msg="p_hat = 0.75: (r - 0.75) / 0.75")
    # Equivalent form from the paper: A_i * (1/N) == r_i / K - 1/N  (K = number of successes)
    N, K = 4, 1
    np.testing.assert_allclose(a[0] / N, r[0] / K - 1 / N, err_msg="should match r_i/K - 1/N after dividing by N")