"""Tests for the Chapter 3 lab. Run:  python -m pytest -q
(With RL_SOLUTIONS=1 the tests run against solutions/mdp.py instead.)"""
import importlib
import os

import numpy as np
import pytest

import envs
from envs import START, N_ACTIONS, N_STATES, state_index

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

UP, RIGHT, DOWN, LEFT = 0, 1, 2, 3


# ---------------------------------------------------------------- Task 1 ---
def test_gridworld_shapes_and_probabilities():
    P, R = mdp.build_gridworld(slip=0.2, step_reward=-0.04)
    assert P.shape == (N_STATES, N_ACTIONS, N_STATES), f"P should have shape (12, 4, 12), got {P.shape}"
    assert R.shape == (N_STATES, N_ACTIONS), f"R should have shape (12, 4), got {R.shape}"
    assert np.all(P >= 0), "probabilities must be non-negative"
    np.testing.assert_allclose(P.sum(axis=2), 1.0, err_msg="every P[s, a, :] must sum to 1")


def test_gridworld_deterministic_moves():
    P, R = mdp.build_gridworld(slip=0.0, step_reward=-0.04)
    s = START  # bottom-left corner (2, 0)
    assert P[s, UP, state_index(1, 0)] == pytest.approx(1.0), "up from the start should reach (1, 0)"
    assert P[s, LEFT, s] == pytest.approx(1.0), "left from the start bumps the edge: you should stay put"
    s = state_index(1, 0)
    assert P[s, RIGHT, s] == pytest.approx(1.0), "right from (1, 0) bumps the wall: you should stay put"
    assert R[START, UP] == pytest.approx(-0.04), "an ordinary step should cost exactly step_reward"
    s = state_index(0, 2)
    assert R[s, RIGHT] == pytest.approx(-0.04 + 1.0), "stepping into the goal earns step_reward + 1"
    s = state_index(2, 3)
    assert R[s, UP] == pytest.approx(-0.04 - 1.0), "stepping into the pit earns step_reward - 1"


def test_gridworld_slip():
    P, R = mdp.build_gridworld(slip=0.2, step_reward=-0.04)
    s = state_index(0, 2)  # just left of the goal, aiming right
    expected = {state_index(0, 3): 0.8, s: 0.1, state_index(1, 2): 0.1}
    for s2, p in expected.items():
        assert P[s, RIGHT, s2] == pytest.approx(p), (
            f"from (0, 2) going right with slip 0.2: P(land in {envs.cell_of(s2)}) should be {p}, got {P[s, RIGHT, s2]}")
    assert R[s, RIGHT] == pytest.approx(-0.04 + 0.8 * 1.0), "R is the EXPECTED reward: -0.04 + 0.8 * (+1)"
    # Two slips that both bump add up: from the top-left corner going up,
    # intended (bump) 0.8 and slip-left (bump) 0.1 both land on the same cell.
    s = state_index(0, 0)
    assert P[s, UP, s] == pytest.approx(0.9), "two bumps into the same cell must ADD (0.8 + 0.1); use +="
    assert P[s, UP, state_index(0, 1)] == pytest.approx(0.1)


def test_gridworld_terminals_are_absorbing():
    P, R = mdp.build_gridworld(slip=0.2, step_reward=-0.04)
    for s in [state_index(0, 3), state_index(1, 3), state_index(1, 1)]:
        for a in range(N_ACTIONS):
            assert P[s, a, s] == pytest.approx(1.0), f"terminal/wall state {envs.cell_of(s)} must loop to itself"
            assert R[s, a] == 0.0, f"terminal/wall state {envs.cell_of(s)} must give reward 0"


# ---------------------------------------------------------------- Task 2 ---
def test_policy_matrices_machine():
    P, R = envs.machine_mdp()
    P_pi, r_pi = mdp.policy_matrices(P, R, envs.MACHINE_RUN_TO_FAILURE)
    np.testing.assert_allclose(P_pi, [[0.6, 0.4, 0.0], [0.0, 0.5, 0.5], [1.0, 0.0, 0.0]],
                               err_msg="P_pi for run-to-failure (run, run, repair) is wrong")
    np.testing.assert_allclose(r_pi, [2.0, 0.0, -10.0], err_msg="r_pi for run-to-failure is wrong")
    # a 50/50 mixed policy averages rows
    pi = np.full((3, 2), 0.5)
    P_pi, r_pi = mdp.policy_matrices(P, R, pi)
    np.testing.assert_allclose(r_pi, [0.5, -0.5, -5.0], err_msg="r_pi must average R over pi")
    np.testing.assert_allclose(P_pi.sum(axis=1), 1.0, err_msg="rows of P_pi must sum to 1")


# ---------------------------------------------------------------- Task 3 ---
def test_exact_machine_matches_hand_solution():
    P, R = envs.machine_mdp()
    V = mdp.evaluate_exact(P, R, envs.MACHINE_RUN_TO_FAILURE, gamma=0.5)
    np.testing.assert_allclose(V, [2.0, -3.0, -9.0], atol=1e-9,
                               err_msg="run-to-failure at gamma=0.5 should give the chapter's (2, -3, -9)")
    V = mdp.evaluate_exact(P, R, envs.MACHINE_MAINTAIN, gamma=0.5)
    np.testing.assert_allclose(V, [3.0, 0.5, -8.5], atol=1e-9,
                               err_msg="maintain at gamma=0.5 should give the chapter's (3, 0.5, -8.5)")


def test_exact_gamma_zero_is_immediate_reward():
    P, R = mdp.build_gridworld()
    pi = envs.uniform_policy(N_STATES, N_ACTIONS)
    V = mdp.evaluate_exact(P, R, pi, gamma=0.0)
    np.testing.assert_allclose(V, (pi * R).sum(axis=1), err_msg="with gamma=0, V^pi(s) must equal r_pi(s)")


def test_exact_satisfies_bellman_and_matches_monte_carlo():
    P, R = mdp.build_gridworld(slip=0.2, step_reward=-0.04)
    pi, gamma = envs.SHORTEST_PATH, 0.9
    V = mdp.evaluate_exact(P, R, pi, gamma)
    # Bellman expectation equation, written out with loops
    for s in range(N_STATES):
        rhs = sum(pi[s, a] * (R[s, a] + gamma * P[s, a] @ V) for a in range(N_ACTIONS))
        assert V[s] == pytest.approx(rhs, abs=1e-9), f"V does not satisfy the Bellman equation at {envs.cell_of(s)}"
    assert V[START] == pytest.approx(0.2392, abs=1e-3), "V(start) for the shortest-path preset should be about 0.239"
    # ... and it really is an average return: sample episodes and compare.
    mc_mean, mc_se = envs.monte_carlo_value(P, R, pi, gamma, START, n_episodes=3000, seed=0)
    assert abs(mc_mean - V[START]) < 4 * mc_se + 1e-3, (
        f"Monte Carlo average return {mc_mean:.3f} ± {mc_se:.3f} disagrees with V(start) = {V[START]:.3f}")


# ---------------------------------------------------------------- Task 4 ---
@pytest.mark.parametrize("gamma", [0.5, 0.9, 0.99])
def test_iterative_matches_exact(gamma):
    P, R = mdp.build_gridworld(slip=0.2, step_reward=-0.04)
    pi = envs.uniform_policy(N_STATES, N_ACTIONS)
    V_exact = mdp.evaluate_exact(P, R, pi, gamma)
    V_iter, n = mdp.evaluate_iterative(P, R, pi, gamma, tol=1e-10)
    np.testing.assert_allclose(V_iter, V_exact, atol=1e-7,
                               err_msg=f"iterative and exact evaluation disagree at gamma={gamma}")
    assert isinstance(n, (int, np.integer)) and n > 1, "evaluate_iterative must return (V, number_of_updates)"


def test_iterative_counts_and_starts_from_zero():
    P, R = envs.machine_mdp()
    pi = envs.MACHINE_MAINTAIN
    # With max_iters=1 you get exactly one update from V_0 = 0, i.e. V_1 = r_pi.
    V1, n1 = mdp.evaluate_iterative(P, R, pi, gamma=0.5, tol=1e-10, max_iters=1)
    np.testing.assert_allclose(V1, [2.0, -1.0, -10.0], err_msg="one update from V_0 = 0 should give V_1 = r_pi")
    assert n1 == 1
    # Error shrinks by about gamma per sweep, so ~ log(1e-10)/log(0.5) = 33 sweeps (give or take).
    _, n = mdp.evaluate_iterative(P, R, pi, gamma=0.5, tol=1e-10)
    assert 20 <= n <= 60, f"expected roughly 30-40 updates at gamma=0.5, tol=1e-10; got {n}"


# ---------------------------------------------------------------- Task 5 ---
def test_q_from_v_machine():
    P, R = envs.machine_mdp()
    V = mdp.evaluate_exact(P, R, envs.MACHINE_RUN_TO_FAILURE, gamma=0.5)
    Q = mdp.q_from_v(P, R, V, gamma=0.5)
    assert Q.shape == (3, 2), f"Q should have shape (3, 2), got {Q.shape}"
    assert Q[1, 0] == pytest.approx(-3.0), "Q(worn, run) under run-to-failure equals V(worn) = -3"
    assert Q[1, 1] == pytest.approx(0.0), "Q(worn, repair) = -1 + 0.5 * V(new) = 0 (the chapter's example)"


def test_v_is_policy_average_of_q():
    P, R = mdp.build_gridworld(slip=0.3, step_reward=-0.1)
    rng = np.random.default_rng(0)
    pi = rng.dirichlet(np.ones(N_ACTIONS), size=N_STATES)  # a random stochastic policy
    V = mdp.evaluate_exact(P, R, pi, gamma=0.95)
    Q = mdp.q_from_v(P, R, V, gamma=0.95)
    np.testing.assert_allclose((pi * Q).sum(axis=1), V, atol=1e-9,
                               err_msg="sum_a pi(a|s) Q(s,a) must equal V(s)")


def test_one_step_lookahead_improves_the_start():
    """'Does it learn' in miniature: acting greedily w.r.t. Q^pi at one state
    can only help (a preview of policy improvement, Chapter 4)."""
    P, R = mdp.build_gridworld(slip=0.2, step_reward=-0.04)
    pi = envs.SHORTEST_PATH
    V = mdp.evaluate_exact(P, R, pi, 0.9)
    Q = mdp.q_from_v(P, R, V, 0.9)
    best = int(np.argmax(Q[START]))
    assert best == UP, "under the shortest-path policy, Q(start, up) should be the largest"
    pi2 = pi.copy()
    pi2[START] = np.eye(N_ACTIONS)[best]
    V2 = mdp.evaluate_exact(P, R, pi2, 0.9)
    assert np.all(V2 >= V - 1e-12) and V2[START] > V[START] + 0.05, "switching the start to 'up' should raise V(start)"