"""Tests for the Chapter 2 lab. Run: python -m pytest -q  (or -k task2 for one task)."""

import importlib
import math
import os

import numpy as np
import pytest

from envs import BernoulliBandit, bernoulli_testbed, run_bandit, run_testbed

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


# ---------------------------------------------------------------- Task 1
def test_task1_unique_max():
    rng = np.random.default_rng(0)
    assert bandits.argmax_tiebreak([0.1, 0.7, 0.3], rng) == 1, "the max is at index 1"
    assert bandits.argmax_tiebreak([-5.0, -1.0, -3.0], rng) == 1, "works with negative numbers too"
    assert bandits.argmax_tiebreak([1.0, np.inf, 2.0], rng) == 1, "inf is the largest value"


def test_task1_ties_are_random():
    rng = np.random.default_rng(1)
    picks = [bandits.argmax_tiebreak([0.0, 2.0, 1.0, 2.0, 2.0], rng) for _ in range(3000)]
    counts = np.bincount(picks, minlength=5)
    assert counts[0] == 0 and counts[2] == 0, f"only tied maxima (1, 3, 4) may be returned, got counts {counts}"
    freqs = counts[[1, 3, 4]] / 3000
    assert np.all(np.abs(freqs - 1 / 3) < 0.04), (
        f"each of the 3 tied arms should be picked about 1/3 of the time, got {freqs.round(3)}")


# ---------------------------------------------------------------- Task 2
def test_task2_sample_average_matches_chapter():
    # The chapter's worked example: rewards 1, 0, 0, 1, 1 from one arm.
    agent = bandits.EpsilonGreedy(3, 0.0, np.random.default_rng(0))
    expected = [1.0, 0.5, 1 / 3, 0.5, 0.6]
    for r, q in zip([1, 0, 0, 1, 1], expected):
        agent.update(2, r)
        assert agent.Q[2] == pytest.approx(q), f"after this reward Q[2] should be {q:.4f}, got {agent.Q[2]:.4f}"
    assert list(agent.N) == [0, 0, 5], f"N should count pulls per arm, expected [0, 0, 5], got {list(agent.N)}"
    assert list(agent.Q[:2]) == [0.0, 0.0], "only the pulled arm's estimate may change"


def test_task2_constant_step_size():
    agent = bandits.EpsilonGreedy(2, 0.0, np.random.default_rng(0), initial_value=0.0, step_size=0.1)
    agent.update(0, 1.0)
    assert agent.Q[0] == pytest.approx(0.1), "Q <- 0 + 0.1 * (1 - 0) = 0.1"
    agent.update(0, 1.0)
    assert agent.Q[0] == pytest.approx(0.19), "Q <- 0.1 + 0.1 * (1 - 0.1) = 0.19"
    agent.update(0, 0.0)
    assert agent.Q[0] == pytest.approx(0.171), "Q <- 0.19 + 0.1 * (0 - 0.19) = 0.171"


def test_task2_optimistic_initial_value():
    agent = bandits.EpsilonGreedy(2, 0.0, np.random.default_rng(0), initial_value=5.0, step_size=0.1)
    assert list(agent.Q) == [5.0, 5.0], "Q should start at initial_value"
    agent.update(1, 0.0)
    assert agent.Q[1] == pytest.approx(4.5), "Q <- 5 + 0.1 * (0 - 5) = 4.5"


def test_task2_greedy_picks_the_max():
    agent = bandits.EpsilonGreedy(4, 0.0, np.random.default_rng(0))
    agent.Q[:] = [0.2, 0.9, 0.5, 0.1]
    picks = {agent.select_action() for _ in range(200)}
    assert picks == {1}, f"with epsilon = 0 the agent must always pick arm 1 (largest Q), got {picks}"


def test_task2_epsilon_explores_at_the_right_rate():
    agent = bandits.EpsilonGreedy(10, 0.2, np.random.default_rng(3))
    agent.Q[:] = np.arange(10) / 10.0        # arm 9 is greedy
    picks = np.array([agent.select_action() for _ in range(20000)])
    p_greedy = np.mean(picks == 9)
    # Greedy arm: (1 - eps) + eps / K = 0.8 + 0.02 = 0.82; every other arm: eps / K = 0.02.
    assert abs(p_greedy - 0.82) < 0.015, (
        f"with eps = 0.2 and K = 10 the greedy arm should be picked with prob 0.82, got {p_greedy:.3f}. "
        "Remember exploration picks uniformly among ALL arms, including the greedy one.")
    p_arm0 = np.mean(picks == 0)
    assert abs(p_arm0 - 0.02) < 0.006, f"each non-greedy arm should be picked with prob 0.02, got {p_arm0:.3f}"


# ---------------------------------------------------------------- Task 3
def test_task3_ucb_scores_match_chapter():
    # The chapter's worked example: t = 20, N = (12, 6, 2), Q = (7/12, 3/6, 0/2), c = sqrt(2).
    s = bandits.ucb_scores([7 / 12, 0.5, 0.0], [12, 6, 2], 20, math.sqrt(2))
    expected = [7 / 12 + math.sqrt(2 * math.log(20) / 12),
                0.5 + math.sqrt(2 * math.log(20) / 6),
                0.0 + math.sqrt(2 * math.log(20) / 2)]
    assert np.allclose(s, expected), f"expected about [1.29, 1.50, 1.73], got {np.round(s, 3)}"


def test_task3_ucb_untried_is_inf():
    with np.errstate(all="raise"):           # dividing by N = 0 would raise here
        s = bandits.ucb_scores([0.5, 0.0, 0.2], [3, 0, 1], 4, 1.0)
    assert math.isinf(s[1]) and s[1] > 0, f"an untried arm must score +inf, got {s[1]}"
    assert np.isfinite(s[0]) and np.isfinite(s[2]), "tried arms must have finite scores"


def test_task3_ucb1_tries_every_arm_first():
    agent = bandits.UCB1(5, np.random.default_rng(0))
    first = []
    for _ in range(5):
        a = agent.select_action()
        first.append(a)
        agent.update(a, 0.0)
    assert sorted(first) == [0, 1, 2, 3, 4], f"the first K pulls must try each arm once, got {first}"
    assert agent.t == 5 and list(agent.N) == [1] * 5, "t and N must count the pulls"


def test_task3_ucb1_follows_the_scores():
    agent = bandits.UCB1(3, np.random.default_rng(0), c=math.sqrt(2))
    agent.Q[:] = [7 / 12, 0.5, 0.0]
    agent.N[:] = [12, 6, 2]
    agent.t = 20
    assert agent.select_action() == 2, "the chapter's example: arm 3 (index 2) has the highest UCB score"
    agent.update(2, 0.0)                       # it fails again
    assert agent.t == 21 and agent.N[2] == 3 and agent.Q[2] == 0.0
    assert agent.select_action() == 1, "after another failure of arm 3, arm 2 (index 1) has the top score"


# ---------------------------------------------------------------- Task 4
def test_task4_beta_update():
    agent = bandits.ThompsonBernoulli(3, np.random.default_rng(0))
    for a, r in [(0, 1), (0, 1), (0, 0), (2, 0)]:
        agent.update(a, r)
    assert list(agent.alpha) == [3.0, 1.0, 1.0], f"alpha = 1 + wins, expected [3, 1, 1], got {list(agent.alpha)}"
    assert list(agent.beta) == [2.0, 1.0, 2.0], f"beta = 1 + losses, expected [2, 1, 2], got {list(agent.beta)}"


def test_task4_probability_matching():
    # Exercise 2.4: arm 0 has Beta(2, 1), arm 1 has Beta(1, 1).
    # Thompson sampling picks arm 0 with probability P(theta_0 > theta_1) = 2/3.
    agent = bandits.ThompsonBernoulli(2, np.random.default_rng(7))
    agent.update(0, 1.0)
    picks = np.array([agent.select_action() for _ in range(6000)])
    p0 = np.mean(picks == 0)
    assert abs(p0 - 2 / 3) < 0.025, f"arm 0 should be picked about 2/3 of the time, got {p0:.3f}"


def test_task4_confident_posterior_exploits():
    agent = bandits.ThompsonBernoulli(2, np.random.default_rng(0))
    agent.alpha[:] = [801, 201]               # arm 0: about 0.8, arm 1: about 0.2, both well measured
    agent.beta[:] = [201, 801]
    picks = {agent.select_action() for _ in range(300)}
    assert picks == {0}, "with sharp posteriors far apart, Thompson sampling should always pick arm 0"


# ---------------------------------------------------------------- Does it learn?
def _final(make_agent, runs=20, steps=1000):
    return run_testbed(make_agent, bernoulli_testbed(10), runs=runs, steps=steps, seed=123)


def test_all_agents_learn():
    rand = _final(lambda k, rng: bandits.EpsilonGreedy(k, 1.0, rng))            # pure exploration
    eps = _final(lambda k, rng: bandits.EpsilonGreedy(k, 0.1, rng))
    ucb = _final(lambda k, rng: bandits.UCB1(k, rng))
    ts = _final(lambda k, rng: bandits.ThompsonBernoulli(k, rng))
    r_rand = rand["cum_regret"][-1]
    for name, out in [("eps-greedy 0.1", eps), ("UCB1", ucb), ("Thompson", ts)]:
        assert out["cum_regret"][-1] < 0.6 * r_rand, (
            f"{name}: regret after 1000 steps is {out['cum_regret'][-1]:.0f}, "
            f"not much better than random play ({r_rand:.0f})")
        assert out["pct_optimal"][-200:].mean() > 0.45, (
            f"{name}: should pick the best arm often by the end, got {out['pct_optimal'][-200:].mean():.2f}")
    assert ts["cum_regret"][-1] < eps["cum_regret"][-1], "Thompson sampling should beat eps-greedy here"


def test_greedy_can_lock_in():
    # The chapter's lock-in story: arms with q = (0.9, 0.5). Greedy (Q starts at 0) gets stuck on the
    # worse arm in a sizable fraction of runs (Exercise 2.2: probability 5/14 ≈ 0.36).
    stuck = 0
    runs = 400
    for i in range(runs):
        rng = np.random.default_rng([99, i])
        agent = bandits.EpsilonGreedy(2, 0.0, rng)
        out = run_bandit(agent, BernoulliBandit([0.9, 0.5], rng), 100)
        stuck += out["optimal"][-50:].mean() == 0.0
    frac = stuck / runs
    assert 0.27 < frac < 0.45, f"greedy should lock onto the 0.5 arm in about 36% of runs, got {frac:.2f}"