import numpy as np
import pytest

from bandits.algorithms import (KLUCB, UCB1, BernoulliEnv, GaussianEnv, ThompsonBeta,
                                ThompsonGaussian, kl_bernoulli, klucb_index,
                                klucb_index_newton, simulate)
from bandits.lower_bound import lai_robbins_constant


@pytest.mark.parametrize("mu,n,level", [(0.1, 10, 2.0), (0.5, 100, 5.0), (0.9, 3, 1.0),
                                        (0.0, 20, 3.0), (0.97, 1000, 9.0)])
def test_klucb_index_bisection_matches_newton_and_definition(mu, n, level):
    q = float(klucb_index(np.array([mu]), np.array([n]), level)[0])
    q_newton = klucb_index_newton(mu, n, level)
    assert abs(q - q_newton) < 1e-7
    assert mu <= q <= 1
    # q is the largest feasible point: feasible at q, infeasible just above
    assert n * kl_bernoulli(mu, q) <= level + 1e-9
    if q < 1 - 1e-6:
        assert n * kl_bernoulli(mu, q + 1e-6) > level


def test_klucb_index_shrinks_with_more_samples():
    qs = klucb_index(np.full(4, 0.3), np.array([1, 10, 100, 1000]), 3.0)
    assert np.all(np.diff(qs) < 0) and qs[-1] > 0.3


def test_ts_beta_posterior_update():
    pol = ThompsonBeta()
    pol.reset(2, 3, np.random.default_rng(0))
    pol.update(np.array([0, 2]), np.array([1.0, 0.0]))
    pol.update(np.array([0, 2]), np.array([1.0, 1.0]))
    np.testing.assert_array_equal(pol.alpha, [[3, 1, 1], [1, 1, 2]])
    np.testing.assert_array_equal(pol.beta, [[1, 1, 1], [1, 1, 2]])


def test_ts_gaussian_posterior_update():
    pol = ThompsonGaussian(sigma=2.0, mu0=0.0, tau0=1.0)
    pol.reset(1, 2, np.random.default_rng(0))
    for r in (1.0, 3.0):
        pol.update(np.array([1]), np.array([r]))
    mean, sd = pol.posterior()
    # precision = 1 + 2/4 = 1.5 ; mean = (4/4)/1.5
    assert mean[0, 1] == pytest.approx(1.0 / 1.5)
    assert sd[0, 1] == pytest.approx(1 / np.sqrt(1.5))
    assert mean[0, 0] == 0 and sd[0, 0] == pytest.approx(1.0)


def test_ucb_regret_sublinear():
    env = BernoulliEnv([0.7, 0.5, 0.4])
    reg = simulate(env, UCB1(), T=4000, n_runs=30, seed=1)["regret"].mean(axis=0)
    # sublinear: second half adds much less regret than the first half
    first, second = reg[1999], reg[-1] - reg[1999]
    assert second < 0.5 * first
    # and far below the linear regret of uniform play
    assert reg[-1] < 0.25 * 4000 * env.gaps.mean()


def test_klucb_and_ts_beat_ucb1_and_identify_best_arm():
    env = BernoulliEnv([0.8, 0.6, 0.5])
    res = {p.name: simulate(env, p, T=3000, n_runs=30, seed=2)
           for p in (UCB1(), KLUCB(), ThompsonBeta())}
    final = {k: v["regret"][:, -1].mean() for k, v in res.items()}
    assert final["KL-UCB"] < final["UCB1"] and final["TS-Beta"] < final["UCB1"]
    for v in res.values():
        assert (v["counts"].argmax(axis=1) == 0).mean() > 0.95


def test_gaussian_env_and_lower_bound_constants():
    env = GaussianEnv([1.0, 0.5, 0.0], sigma=1.0)
    assert lai_robbins_constant(env.means, "gaussian", 1.0) == pytest.approx(2 / 0.5 + 2 / 1.0)
    means = [0.9, 0.8]
    assert lai_robbins_constant(means) == pytest.approx(0.1 / kl_bernoulli(0.8, 0.9))
    r = simulate(env, ThompsonGaussian(), T=2000, n_runs=20, seed=3)["regret"][:, -1]
    assert r.mean() < 100