import numpy as np

from bandits.multiplayer import (SICSimplified, collisions_and_regret,
                                 estimate_num_players, musical_chairs, selfish_ucb)

MEANS = np.array([0.9, 0.8, 0.7, 0.5, 0.4, 0.3])


def test_collision_rules():
    arms = np.array([[0, 0, 1],   # players 0,1 collide on arm 0
                     [0, 1, 2],   # optimal, no collision
                     [3, 3, 3]])  # everyone collides
    collided, regret = collisions_and_regret(arms, MEANS, 3)
    np.testing.assert_array_equal(collided, [[1, 1, 0], [0, 0, 0], [1, 1, 1]])
    best = 0.9 + 0.8 + 0.7
    np.testing.assert_allclose(regret, [best - 0.8, 0.0, best])


def test_collided_players_get_zero_reward():
    # with every player forced onto one arm (K=1 effectively), rewards are 0
    res = selfish_ucb(np.array([1.0]), M=2, T=50, n_runs=3, seed=0)
    assert np.all(res["collisions"] == 2 * 50)
    np.testing.assert_allclose(res["regret"][:, -1], 50 * 1.0)


def test_estimate_num_players_inverts_collision_rate():
    K, T0 = 6, 100000
    for M in (1, 2, 3, 5):
        p_coll = 1 - (1 - 1 / K) ** (M - 1)
        assert estimate_num_players(p_coll * T0, T0, K) == M


def test_musical_chairs_settles_on_top_m():
    res = musical_chairs(MEANS, M=3, T=6000, T0=1500, n_runs=10, seed=0)
    last = res["regret"][:, -1] - res["regret"][:, -2]
    assert np.mean(last < 1e-9) >= 0.9


def test_sic_simplified_orthogonal_and_logarithmic():
    sic = SICSimplified()
    r1 = sic.run(MEANS, 3, 5000, n_runs=10, seed=0)
    r2 = sic.run(MEANS, 3, 50000, n_runs=10, seed=0)
    assert r1["failures"] == 0 and r2["failures"] == 0
    last = r2["regret"][:, -1] - r2["regret"][:, -2]
    assert np.all(last < 1e-9)  # committed to the true top-3 set
    # 10x the horizon costs far less than 10x the regret
    assert r2["regret"][:, -1].mean() < 2.0 * r1["regret"][:, -1].mean()


def test_sic_single_player_works():
    r = SICSimplified().run(MEANS, 1, 20000, n_runs=5, seed=1)
    assert np.all(r["regret"][:, -1] - r["regret"][:, -2] < 1e-9)