"""Tests for the Chapter 4 lab. Run:  python -m pytest -q   (or -k task2 for one task)."""
import importlib
import itertools
import os

import numpy as np
import pytest

from envs import (DEMO_GRID, FROZEN_LAKE_4x4, FROZEN_LAKE_8x8, deterministic_policy,
                  make_corridor, make_frozen_lake, make_gridworld, random_policy)

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

LEFT, RIGHT = 0, 1


def exact_v(P, R, policy, gamma):
    """V^pi by a linear solve (Chapter 3): V = (I - gamma P_pi)^-1 r_pi."""
    P_pi = np.einsum("sa,sat->st", policy, P)
    r_pi = (policy * R).sum(axis=1)
    return np.linalg.solve(np.eye(len(r_pi)) - gamma * P_pi, r_pi)


# --------------------------------------------------------------------- Task 1
def test_task1_q_from_v_corridor():
    m = make_corridor()
    V = np.array([0.81, 0.9, 1.0, 0.0])
    Q = dp.q_from_v(m.P, m.R, V, 0.9)
    expected = np.array([[0.5, 0.81], [0.729, 0.9], [0.81, 1.0], [0.0, 0.0]])
    assert Q.shape == (4, 2), f"Q should have shape (S, A) = (4, 2), got {Q.shape}"
    assert np.allclose(Q, expected), f"corridor Q with gamma=0.9 should be\n{expected}\ngot\n{Q}"


def test_task1_q_from_v_stochastic():
    m = make_gridworld(DEMO_GRID, slip=0.2, step_reward=-0.04)
    rng = np.random.default_rng(0)
    V = rng.normal(size=m.n_states)
    Q = dp.q_from_v(m.P, m.R, V, 0.95)
    expected = np.array([[m.R[s, a] + 0.95 * sum(m.P[s, a, t] * V[t] for t in range(m.n_states))
                          for a in range(4)] for s in range(m.n_states)])
    assert np.allclose(Q, expected), "q_from_v disagrees with the explicit double sum on a slippery grid"


# --------------------------------------------------------------------- Task 2
def test_task2_policy_evaluation_by_hand():
    # "Always left" in the corridor, gamma = 0.9: V = (0.5, 0.45, 0.405), exact after 3 sweeps.
    m = make_corridor()
    pi = deterministic_policy([LEFT, LEFT, LEFT, LEFT], 2)
    V, n = dp.policy_evaluation(m.P, m.R, pi, 0.9)
    assert np.allclose(V, [0.5, 0.45, 0.405, 0.0]), f"expected V = [0.5, 0.45, 0.405, 0], got {V}"
    assert n <= 5, f"a 3-state deterministic chain needs 3 sweeps plus one to confirm, you used {n}"


def test_task2_policy_evaluation_matches_linear_solve():
    m = make_frozen_lake(FROZEN_LAKE_4x4)
    pi = random_policy(m.n_states, m.n_actions)
    V, _ = dp.policy_evaluation(m.P, m.R, pi, 0.95, tol=1e-12)
    assert np.allclose(V, exact_v(m.P, m.R, pi, 0.95), atol=1e-8), \
        "iterative evaluation of the random policy on FrozenLake should match the linear solve"


def test_task2_policy_evaluation_is_synchronous():
    # With in-place (Gauss-Seidel) updates, sweep 1 would already propagate past state A.
    m = make_corridor()
    pi = deterministic_policy([LEFT, LEFT, LEFT, LEFT], 2)
    V1, n = dp.policy_evaluation(m.P, m.R, pi, 0.9, max_sweeps=1)
    assert n == 1
    assert np.allclose(V1, [0.5, 0.0, 0.0, 0.0]), \
        f"after ONE synchronous sweep from zeros, only A knows about its reward: expected [0.5, 0, 0, 0], got {V1}"


# --------------------------------------------------------------------- Task 3
def test_task3_greedy_policy_corridor():
    m = make_corridor()
    pi = dp.greedy_policy(m.P, m.R, np.array([0.81, 0.9, 1.0, 0.0]), 0.9)
    assert pi.shape == (4, 2)
    assert np.allclose(pi.sum(axis=1), 1) and set(np.unique(pi)) <= {0.0, 1.0}, "rows must be one-hot"
    assert list(np.argmax(pi[:3], axis=1)) == [RIGHT, RIGHT, RIGHT], "greedy w.r.t. V* should go right everywhere"


def test_task3_greedy_policy_tie_breaking_and_myopia():
    m = make_corridor()
    pi = dp.greedy_policy(m.P, m.R, np.zeros(4), 0.9)
    # With V = 0 only immediate rewards matter: A prefers the +0.5 exit, B is a tie (-> action 0), C exits right.
    assert list(np.argmax(pi, axis=1)) == [LEFT, LEFT, RIGHT, LEFT], \
        f"with V = 0: expected actions [left, left (tie), right, left (tie)], got {list(np.argmax(pi, axis=1))}"


# --------------------------------------------------------------------- Task 4
def test_task4_value_iteration_corridor_by_hand():
    m = make_corridor()
    V, pi, deltas = dp.value_iteration(m.P, m.R, 0.9)
    assert np.allclose(V, [0.81, 0.9, 1.0, 0.0]), f"V* for gamma = 0.9 is [0.81, 0.9, 1, 0], got {V}"
    assert np.allclose(deltas, [1.0, 0.9, 0.31, 0.0]), \
        f"the chapter's table gives max changes [1.0, 0.9, 0.31, 0.0] per sweep, got {deltas}"
    assert list(np.argmax(pi[:3], axis=1)) == [RIGHT, RIGHT, RIGHT]


def test_task4_value_iteration_small_gamma_changes_policy():
    m = make_corridor()
    V, pi, _ = dp.value_iteration(m.P, m.R, 0.5)
    assert np.allclose(V, [0.5, 0.5, 1.0, 0.0]), f"V* for gamma = 0.5 is [0.5, 0.5, 1, 0], got {V}"
    assert np.argmax(pi[0]) == LEFT, "with gamma = 0.5, A should take the small exit now"


def test_task4_value_iteration_contracts():
    m = make_gridworld(DEMO_GRID, slip=0.2)
    gamma = 0.9
    _, _, deltas = dp.value_iteration(m.P, m.R, gamma, tol=1e-9)
    d = np.array(deltas)
    assert np.all(d[1:] <= gamma * d[:-1] + 1e-12), \
        "each sweep must shrink the max change by at least a factor gamma (the contraction property)"


def test_task4_value_iteration_any_start():
    m = make_gridworld(DEMO_GRID, slip=0.1)
    V_a, _, _ = dp.value_iteration(m.P, m.R, 0.9, tol=1e-12)
    rng = np.random.default_rng(1)
    V0 = rng.uniform(-5, 5, size=m.n_states)
    V_b, _, _ = dp.value_iteration(m.P, m.R, 0.9, tol=1e-12, V0=V0)
    assert np.allclose(V_a, V_b, atol=1e-8), "value iteration must reach the same V* from any starting V"


# --------------------------------------------------------------------- Task 5
def test_task5_policy_iteration_corridor_by_hand():
    m = make_corridor()
    pi0 = deterministic_policy([LEFT, LEFT, LEFT, LEFT], 2)
    pi, V, history = dp.policy_iteration(m.P, m.R, 0.9, policy0=pi0)
    seq = [tuple(np.argmax(h[:3], axis=1)) for h in history]
    expected = [(0, 0, 0), (0, 0, 1), (0, 1, 1), (1, 1, 1)]
    assert seq == expected, f"policies should go LLL -> LLR -> LRR -> RRR (0 = L, 1 = R), got {seq}"
    assert np.allclose(V, [0.81, 0.9, 1.0, 0.0])


def test_task5_policy_iteration_agrees_with_value_iteration():
    m = make_frozen_lake(FROZEN_LAKE_8x8)
    gamma = 0.99
    pi, V_pi, history = dp.policy_iteration(m.P, m.R, gamma)
    V_vi, _, _ = dp.value_iteration(m.P, m.R, gamma, tol=1e-12)
    assert np.allclose(V_pi, V_vi, atol=1e-6), "policy iteration and value iteration must find the same V*"
    assert len(history) < 30, f"policy iteration should need few iterations here, used {len(history)}"
    assert np.allclose(exact_v(m.P, m.R, pi, gamma), V_vi, atol=1e-6), "the returned policy must achieve V*"


# ---------------------------------------------------- all tasks: brute force
def test_all_tasks_beat_brute_force():
    """V* is the elementwise max over ALL deterministic policies; check DP finds it."""
    m = make_gridworld(["..G", ".P."], slip=0.2, step_reward=-0.05)
    gamma = 0.9
    free = [s for s in range(m.n_states) if not m.terminal[s]]
    best = np.full(m.n_states, -np.inf)
    for acts in itertools.product(range(4), repeat=len(free)):
        a = np.zeros(m.n_states, dtype=int)
        a[free] = acts
        best = np.maximum(best, exact_v(m.P, m.R, deterministic_policy(a, 4), gamma))
    V_vi, pi_vi, _ = dp.value_iteration(m.P, m.R, gamma, tol=1e-12)
    _, V_pi, _ = dp.policy_iteration(m.P, m.R, gamma)
    assert np.allclose(V_vi, best, atol=1e-8), f"value iteration: expected brute-force V* {best}, got {V_vi}"
    assert np.allclose(V_pi, best, atol=1e-8), f"policy iteration: expected brute-force V* {best}, got {V_pi}"
    assert np.allclose(exact_v(m.P, m.R, pi_vi, gamma), best, atol=1e-8), \
        "the greedy policy from value iteration should be optimal in EVERY state at once"