"""Tests for Part 1 (overopt.py). Run:  python -m pytest -q -k task1"""

import importlib
import math
import os

import numpy as np
import pytest

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

from common import make_pool  # noqa: E402


# ---------------------------------------------------------------- Task 1a
def test_task1a_bon_kl_values():
    assert float(overopt.bon_kl(1)) == pytest.approx(0.0, abs=1e-12), "best-of-1 is the base policy: KL must be 0"
    assert float(overopt.bon_kl(2)) == pytest.approx(math.log(2) - 0.5, abs=1e-9)
    assert float(overopt.bon_kl(16)) == pytest.approx(1.8351, abs=1e-4), "chapter table: n = 16 gives KL ≈ 1.835 nats"
    arr = overopt.bon_kl(np.array([1, 4, 1000]))
    np.testing.assert_allclose(arr, [0.0, math.log(4) - 0.75, math.log(1000) - 0.999], atol=1e-9,
                               err_msg="bon_kl must work elementwise on arrays")


def test_task1a_bon_kl_matches_exact_kl_on_a_continuous_example():
    # For a continuous proxy, the best-of-n density is n F^(n-1) f. Check the formula by
    # numerically integrating KL = E[log(n F^(n-1))] with F uniform on [0, 1].
    n = 7
    u = (np.arange(200000) + 0.5) / 200000
    dens = n * u ** (n - 1)
    kl = np.mean(dens * np.log(dens))
    assert float(overopt.bon_kl(n)) == pytest.approx(kl, abs=1e-3)


# ---------------------------------------------------------------- Task 1b
def test_task1b_hand_example():
    # Pool of 3 with proxies 1, 3, 2 and golds 10, 0, 5. Ranks by proxy: (1,10), (2,5), (3,0).
    # n = 2: w = [1/9, 4/9 - 1/9, 1 - 4/9] = [1/9, 3/9, 5/9]
    proxy = np.array([1.0, 3.0, 2.0])
    gold = np.array([10.0, 0.0, 5.0])
    p, g = overopt.best_of_n_stats(proxy, gold, 2)
    assert p == pytest.approx((1 * 1 + 2 * 3 + 3 * 5) / 9), "expected proxy (1·1 + 2·3 + 3·5)/9 = 22/9"
    assert g == pytest.approx((10 * 1 + 5 * 3 + 0 * 5) / 9), "expected gold (10·1 + 5·3 + 0·5)/9 = 25/9"


def test_task1b_n1_is_the_pool_mean():
    proxy, gold = make_pool(1000, seed=3)
    p, g = overopt.best_of_n_stats(proxy, gold, 1)
    assert p == pytest.approx(proxy.mean()) and g == pytest.approx(gold.mean()), \
        "best-of-1 is just a random draw: its expectation is the pool mean"


def test_task1b_large_n_picks_the_max():
    proxy, gold = make_pool(500, seed=4)
    p, g = overopt.best_of_n_stats(proxy, gold, 100000)
    i = int(np.argmax(proxy))
    assert p == pytest.approx(proxy[i], abs=1e-6) and g == pytest.approx(gold[i], abs=1e-6)


def test_task1b_matches_simulation():
    proxy, gold = make_pool(2000, b=1.0, seed=5)
    rng = np.random.default_rng(0)
    n = 8
    idx = rng.integers(0, len(proxy), size=(40000, n))
    win = idx[np.arange(len(idx)), np.argmax(proxy[idx], axis=1)]
    p, g = overopt.best_of_n_stats(proxy, gold, n)
    assert p == pytest.approx(proxy[win].mean(), abs=0.02)
    assert g == pytest.approx(gold[win].mean(), abs=0.02)


# ---------------------------------------------------------------- Task 1c
def test_task1c_recovers_exact_coefficients():
    d = np.linspace(0, 3, 20)
    a, b = overopt.fit_gao_bon(d, 1.3 * d - 0.4 * d ** 2)
    assert a == pytest.approx(1.3, abs=1e-8) and b == pytest.approx(0.4, abs=1e-8), \
        f"expected (alpha, beta) = (1.3, 0.4), got ({a}, {b}). Did you use -d**2 as the second column?"


def test_task1c_overoptimization_on_the_toy_pool():
    # With a length-biased proxy, the fitted beta must be positive (gold eventually falls).
    proxy, gold = make_pool(100000, b=1.0, sigma=0.0, seed=0)
    ns = np.unique(np.round(np.logspace(0, 3, 25)).astype(int))
    d = np.sqrt(overopt.bon_kl(ns))
    gain = np.array([overopt.best_of_n_stats(proxy, gold, n)[1] for n in ns]) - gold.mean()
    a, b = overopt.fit_gao_bon(d, gain)
    assert a > 0 and b > 0, f"expected alpha > 0 and beta > 0 for a length-biased proxy, got ({a:.3f}, {b:.3f})"