"""Tests for the Chapter 5 lab. Run with `python -m pytest -q` from this folder.

Set RL_SOLUTIONS=1 to run the same tests against the reference solution.
"""

import importlib
import os

import numpy as np
import pytest

from envs import CliffWalking, RandomWalk, generate_episode, greedy_path, prediction_curve, train

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


# The hand-worked episode from Exercise 5.1 of the chapter:
#   X --(r=1)--> Y --(r=0)--> Z --(r=4)--> end,   states X, Y, Z = 0, 1, 2
EX51_EPISODE = [(0, 1.0, 1, False), (1, 0.0, 2, False), (2, 4.0, None, True)]
EX51_V0 = [1.0, 2.0, 0.0]


# ---------------------------------------------------------------------------
# Task 1: Monte Carlo
# ---------------------------------------------------------------------------

def test_mc_update_hand_example():
    V = np.array(EX51_V0)
    out = td.mc_update(V, EX51_EPISODE, alpha=0.5, gamma=0.9)
    assert out is None, "mc_update should modify V in place and return None"
    # returns: G_2 = 4, G_1 = 0 + 0.9*4 = 3.6, G_0 = 1 + 0.9*3.6 = 4.24
    np.testing.assert_allclose(V, [2.62, 2.8, 2.0], atol=1e-9,
                               err_msg="MC: expected V = [1 + 0.5(4.24-1), 2 + 0.5(3.6-2), 0 + 0.5(4-0)]")


def test_mc_update_random_walk_first_episode():
    # C -> D -> E -> exit right (+1). Every state on the path moves toward G = 1.
    V = np.full(5, 0.5)
    ep = [(2, 0.0, 3, False), (3, 0.0, 4, False), (4, 1.0, None, True)]
    td.mc_update(V, ep, alpha=0.1, gamma=1.0)
    np.testing.assert_allclose(V, [0.5, 0.5, 0.55, 0.55, 0.55], atol=1e-12,
                               err_msg="MC should move C, D and E from 0.5 to 0.55")


def test_mc_update_every_visit():
    # C -> B -> C -> B -> A -> exit left (0). C and B are each visited twice.
    V = np.full(5, 0.5)
    ep = [(2, 0.0, 1, False), (1, 0.0, 2, False), (2, 0.0, 1, False),
          (1, 0.0, 0, False), (0, 0.0, None, True)]
    td.mc_update(V, ep, alpha=0.1, gamma=1.0)
    # each visit multiplies (V - 0) by 0.9
    np.testing.assert_allclose(V, [0.45, 0.405, 0.405, 0.5, 0.5], atol=1e-12,
                               err_msg="every-visit MC: a state visited twice is updated twice")


# ---------------------------------------------------------------------------
# Task 2: TD(0)
# ---------------------------------------------------------------------------

def test_td0_update_hand_example():
    V = np.array(EX51_V0)
    out = td.td0_update(V, EX51_EPISODE, alpha=0.5, gamma=0.9)
    assert out is None, "td0_update should modify V in place and return None"
    # deltas: 1 + 0.9*2 - 1 = 1.8 ; 0 + 0.9*0 - 2 = -2 ; 4 - 0 = 4
    np.testing.assert_allclose(V, [1.9, 1.0, 2.0], atol=1e-9,
                               err_msg="TD(0): expected V = [1 + 0.5*1.8, 2 + 0.5*(-2), 0 + 0.5*4]")


def test_td0_random_walk_two_episodes():
    # The worked example in the chapter: information creeps back one state per episode.
    V = np.full(5, 0.5)
    ep = [(2, 0.0, 3, False), (3, 0.0, 4, False), (4, 1.0, None, True)]
    td.td0_update(V, ep, alpha=0.1, gamma=1.0)
    np.testing.assert_allclose(V, [0.5, 0.5, 0.5, 0.5, 0.55], atol=1e-12,
                               err_msg="after episode 1, TD should change only E (to 0.55)")
    td.td0_update(V, ep, alpha=0.1, gamma=1.0)
    np.testing.assert_allclose(V, [0.5, 0.5, 0.5, 0.505, 0.595], atol=1e-12,
                               err_msg="after episode 2, expected D = 0.505 and E = 0.595")


def test_td0_is_online():
    # A -> B -> A -> exit left. The second update of A must see the new value of B.
    V = np.array([0.0, 1.0, 0.0, 0.0, 0.0])
    ep = [(0, 0.0, 1, False), (1, 0.0, 0, False), (0, 0.0, None, True)]
    td.td0_update(V, ep, alpha=0.5, gamma=1.0)
    # A: 0 + 0.5(1 - 0) = 0.5 ; B: 1 + 0.5(0.5 - 1) = 0.75 ; A: 0.5 + 0.5(0 - 0.5) = 0.25
    np.testing.assert_allclose(V[:2], [0.25, 0.75], atol=1e-12,
                               err_msg="TD(0) updates must be applied one transition at a time")


def test_td0_beats_mc_on_random_walk():
    td_err = prediction_curve(td.td0_update, alpha=0.1, n_episodes=100, n_runs=50)
    mc_err = prediction_curve(td.mc_update, alpha=0.1, n_episodes=100, n_runs=50)
    assert td_err[-1] < 0.08, f"TD(0) RMS error after 100 episodes should be < 0.08, got {td_err[-1]:.3f}"
    assert td_err[-1] < mc_err[-1], (
        f"with alpha = 0.1, TD should end with lower error than MC (TD {td_err[-1]:.3f}, MC {mc_err[-1]:.3f})")


# ---------------------------------------------------------------------------
# Task 3: epsilon-greedy
# ---------------------------------------------------------------------------

def test_epsilon_greedy_greedy_when_epsilon_zero():
    rng = np.random.default_rng(0)
    q = np.array([0.1, 0.7, -0.3, 0.2])
    picks = {td.epsilon_greedy(q, 0.0, rng) for _ in range(200)}
    assert picks == {1}, f"with epsilon = 0 the greedy action 1 should always be picked, got {picks}"
    assert isinstance(td.epsilon_greedy(q, 0.0, rng), (int, np.integer))


def test_epsilon_greedy_breaks_ties_randomly():
    rng = np.random.default_rng(1)
    q = np.array([0.0, 2.0, 2.0, 0.0])
    counts = np.bincount([td.epsilon_greedy(q, 0.0, rng) for _ in range(2000)], minlength=4)
    assert counts[0] == 0 and counts[3] == 0, f"only tied-best actions 1 and 2 may be picked, got {counts}"
    assert 800 < counts[1] < 1200, f"ties should be broken uniformly at random, got counts {counts}"


def test_epsilon_greedy_probabilities():
    rng = np.random.default_rng(2)
    q = np.array([0.0, 1.0, 0.0, 0.0])
    n = 20000
    freq = np.bincount([td.epsilon_greedy(q, 0.2, rng) for _ in range(n)], minlength=4) / n
    # greedy: 1 - 0.2 + 0.2/4 = 0.85 ; each other action: 0.2/4 = 0.05
    np.testing.assert_allclose(freq, [0.05, 0.85, 0.05, 0.05], atol=0.015,
                               err_msg="epsilon = 0.2 with 4 actions: expected frequencies [0.05, 0.85, 0.05, 0.05]")


# ---------------------------------------------------------------------------
# Task 4: the three TD targets
# ---------------------------------------------------------------------------

Q_TARGETS = np.array([[0.0, 0.0, 0.0, 0.0],
                      [1.0, 3.0, 2.0, 0.0]])


def test_sarsa_target():
    got = td.sarsa_target(Q_TARGETS, r=-1.0, s_next=1, a_next=2, done=False, gamma=0.9)
    assert got == pytest.approx(-1 + 0.9 * 2.0), "SARSA target uses Q(s', a') for the action actually taken"
    assert td.sarsa_target(Q_TARGETS, r=-1.0, s_next=1, a_next=2, done=True, gamma=0.9) == pytest.approx(-1.0), \
        "at a terminal transition the target is just r"


def test_q_learning_target():
    got = td.q_learning_target(Q_TARGETS, r=-1.0, s_next=1, done=False, gamma=0.9)
    assert got == pytest.approx(-1 + 0.9 * 3.0), "Q-learning target uses max over a' of Q(s', a') = 3"
    assert td.q_learning_target(Q_TARGETS, r=-1.0, s_next=1, done=True, gamma=0.9) == pytest.approx(-1.0)


def test_expected_sarsa_target():
    # pi(.|s') = [0.05, 0.85, 0.05, 0.05]; E[Q] = 0.05*(1+2+0) + 0.85*3 = 2.7
    got = td.expected_sarsa_target(Q_TARGETS, r=-1.0, s_next=1, done=False, gamma=0.9, epsilon=0.2)
    assert got == pytest.approx(-1 + 0.9 * 2.7), f"expected -1 + 0.9 * 2.7 = 1.43, got {got}"
    # with epsilon = 0 the target policy is greedy and Expected SARSA = Q-learning
    got0 = td.expected_sarsa_target(Q_TARGETS, r=-1.0, s_next=1, done=False, gamma=0.9, epsilon=0.0)
    assert got0 == pytest.approx(-1 + 0.9 * 3.0), "with epsilon = 0, Expected SARSA should equal Q-learning"
    # ties: [2, 2, 0, 0] with epsilon = 0.2 -> 0.8 * 2 + 0.2 * mean = 1.6 + 0.2 = 1.8
    Qt = np.array([[0.0] * 4, [2.0, 2.0, 0.0, 0.0]])
    got_t = td.expected_sarsa_target(Qt, r=0.0, s_next=1, done=False, gamma=1.0, epsilon=0.2)
    assert got_t == pytest.approx(1.8), f"tied maxima: expected 1.8, got {got_t}"
    assert td.expected_sarsa_target(Q_TARGETS, r=-1.0, s_next=1, done=True, gamma=0.9, epsilon=0.2) \
        == pytest.approx(-1.0)


# ---------------------------------------------------------------------------
# Task 5: SARSA and Q-learning episodes
# ---------------------------------------------------------------------------

class Corridor:
    """Two states, 0 and 1. Actions: 0 = left, 1 = right. Every step gives -1.

    Right from state 1 ends the episode; left from state 0 stays in state 0.
    Optimal: go right twice, so Q*(0, right) = -2 and Q*(1, right) = -1.
    Under the uniformly random policy, Q(0, right) = -5 and Q(1, right) = -1.
    """

    n_states, n_actions = 2, 2

    def reset(self):
        self.s = 0
        return 0

    def step(self, a):
        if a == 1 and self.s == 1:
            return 1, -1.0, True
        self.s = min(1, self.s + 1) if a == 1 else max(0, self.s - 1)
        return self.s, -1.0, False


@pytest.mark.parametrize("name", ["sarsa_episode", "q_learning_episode"])
def test_control_exact_episode_with_epsilon_zero(name):
    # With epsilon = 0 and no ties, both methods are deterministic and identical:
    # step 1: Q[0,1] <- -1 + 0.5 * ((-1 + Q[1,1]) - (-1)) = -1 + 0.5 * (-3 + 1) = -2
    # step 2: Q[1,1] <- -2 + 0.5 * (-1 - (-2)) = -1.5
    Q = np.array([[-3.0, -1.0], [-3.0, -2.0]])
    total = getattr(td, name)(Corridor(), Q, 0.5, 0.0, 1.0, np.random.default_rng(0))
    assert total == pytest.approx(-2.0), f"{name} should return the episode's total reward, -2"
    np.testing.assert_allclose(Q, [[-3.0, -2.0], [-3.0, -1.5]], atol=1e-12,
                               err_msg=f"{name}: Q after one greedy episode")


def test_control_off_policy_q_learning_vs_on_policy_sarsa():
    # Behave completely at random (epsilon = 1). Q-learning still learns Q* (off-policy);
    # SARSA learns the values of the random policy it is following (on-policy).
    env = Corridor()
    Q_q = np.zeros((2, 2))
    Q_s = np.zeros((2, 2))
    rng = np.random.default_rng(0)
    for _ in range(4000):
        td.q_learning_episode(env, Q_q, 0.02, 1.0, 1.0, rng)
        td.sarsa_episode(env, Q_s, 0.02, 1.0, 1.0, rng)
    assert Q_q[0, 1] == pytest.approx(-2.0, abs=0.1), f"Q-learning: Q(0, right) should approach Q* = -2, got {Q_q[0, 1]:.2f}"
    assert Q_s[0, 1] == pytest.approx(-5.0, abs=1.2), \
        f"SARSA with epsilon = 1: Q(0, right) should approach the random policy's value -5, got {Q_s[0, 1]:.2f}"


def test_control_cliff_walking_paths_and_returns():
    # Same settings as the chapter's demo: alpha = 0.5, epsilon = 0.1, 500 episodes.
    env = CliffWalking()
    Q_q, _ = train(td.q_learning_episode, env, 500, alpha=0.5, epsilon=0.1, seed=0)
    path_q, ok_q = greedy_path(Q_q)
    assert ok_q and len(path_q) - 1 == 13, \
        f"Q-learning's greedy path should be the 13-step path along the cliff, got {len(path_q) - 1} steps (reached goal: {ok_q})"

    safe = 0
    online_s, online_q = [], []
    for seed in range(3):
        Q_s, ret_s = train(td.sarsa_episode, env, 500, alpha=0.5, epsilon=0.1, seed=seed)
        _, ret_q = train(td.q_learning_episode, env, 500, alpha=0.5, epsilon=0.1, seed=seed)
        path_s, ok_s = greedy_path(Q_s)
        safe += ok_s and len(path_s) - 1 > 13
        online_s.append(ret_s[-200:].mean())
        online_q.append(ret_q[-200:].mean())
    assert safe >= 2, "SARSA's greedy path should usually be a safe detour (> 13 steps) away from the cliff"
    assert np.mean(online_s) > np.mean(online_q) + 8, (
        f"while exploring, SARSA should earn more per episode than Q-learning "
        f"(last 200 episodes: SARSA {np.mean(online_s):.1f}, Q-learning {np.mean(online_q):.1f}; "
        f"expect roughly -25 vs -50)")