import importlib
import os

import numpy as np
import pytest

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

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


# ---------------------------------------------------------------- Task 1
def test_discounted_return_hand_cases():
    assert gw.discounted_return([], 0.9) == 0.0, "an empty episode has return 0"
    assert gw.discounted_return([0, 0, 1], 0.9) == pytest.approx(0.81), "0 + 0.9*0 + 0.81*1"
    assert gw.discounted_return([1, 1, 1, 1], 1.0) == pytest.approx(4.0), "gamma = 1: plain sum"
    assert gw.discounted_return([5, 7, 9], 0.0) == pytest.approx(5.0), "gamma = 0: only r_0 counts"
    # The chapter's worked example: -0.04 per step, +1 on the 6th step, gamma = 0.9
    rewards = [-0.04] * 5 + [1.0]
    expected = sum(-0.04 * 0.9**k for k in range(5)) + 0.9**5
    assert gw.discounted_return(rewards, 0.9) == pytest.approx(expected)
    assert isinstance(gw.discounted_return([1, 2], 0.5), float), "return a Python float"


def test_discounted_return_accepts_arrays():
    assert gw.discounted_return(np.array([1.0, 2.0, 3.0]), 0.5) == pytest.approx(1 + 1 + 0.75)


# ---------------------------------------------------------------- Task 2
def test_reset():
    env = gw.GridWorld(seed=0)
    s = env.reset()
    assert s == (2, 0), f"start state should be (2, 0), got {s}"
    assert env.state == (2, 0) and env.done is False


def test_step_deterministic_moves_walls_and_edges():
    env = gw.GridWorld(slip=0.0, seed=0)
    env.reset()
    assert env.step(LEFT) == ((2, 0), 0.0, False), "moving off the grid leaves you in place"
    assert env.step(UP) == ((1, 0), 0.0, False)
    assert env.step(RIGHT) == ((1, 0), 0.0, False), "(1, 1) is a wall: stay in place"
    assert env.step(UP) == ((0, 0), 0.0, False)
    assert env.step(UP) == ((0, 0), 0.0, False), "top edge: stay in place"
    assert env.step(DOWN) == ((1, 0), 0.0, False)


def test_step_goal_pit_and_step_reward():
    env = gw.GridWorld(slip=0.0, step_reward=-0.04, seed=0)
    env.reset()
    for _ in range(5):
        s, r, done = env.step(RIGHT)
        assert r == pytest.approx(-0.04) and not done, "ordinary moves give step_reward"
    s, r, done = env.step(RIGHT)
    assert (s, r, done) == ((2, 6), 1.0, True), "landing on G: reward +1 and done"
    with pytest.raises(RuntimeError):
        env.step(RIGHT)

    env.reset()
    env.step(RIGHT)
    s, r, done = env.step(DOWN)
    assert (s, r, done) == ((3, 1), -1.0, True), "landing on P: reward -1 and done"


def test_step_slip_statistics():
    env = gw.GridWorld(slip=0.4, seed=123)
    counts = {}
    n = 6000
    for _ in range(n):
        env.reset()
        env.state = (0, 3)  # middle of the top row: up is the edge, down is a wall
        s, _, _ = env.step(LEFT)
        counts[s] = counts.get(s, 0) + 1
    # intended: (0, 2) with prob 0.6. Perpendicular moves (up / down) are blocked: stay at (0, 3).
    assert set(counts) <= {(0, 2), (0, 3)}, f"left from (0,3) can only reach (0,2) or stay; got {counts}"
    frac = counts.get((0, 2), 0) / n
    assert abs(frac - 0.6) < 0.03, f"with slip=0.4 the intended move should happen ~60% of the time, got {frac:.3f}"

    counts = {}
    for _ in range(n):
        env.reset()
        env.state = (2, 3)  # on the bridge, pressing right
        s, r, done = env.step(RIGHT)
        counts[s] = counts.get(s, 0) + 1
    fall = counts.get((3, 3), 0) / n
    stay = counts.get((2, 3), 0) / n
    assert abs(fall - 0.2) < 0.03, f"slip=0.4: should fall into the pit ~20% of the time, got {fall:.3f}"
    assert abs(stay - 0.2) < 0.03, f"slip=0.4: should bump the wall above ~20% of the time, got {stay:.3f}"


def test_step_uses_env_rng():
    a = gw.GridWorld(slip=0.5, seed=7)
    b = gw.GridWorld(slip=0.5, seed=7)
    for env in (a, b):
        env.reset()
    ta = [a.step(UP)[0] if not a.done else None for _ in range(20)]
    tb = [b.step(UP)[0] if not b.done else None for _ in range(20)]
    assert ta == tb, "same seed must give the same trajectory: use self.rng for all randomness"


# ---------------------------------------------------------------- Task 3
def test_random_policy():
    rng = np.random.default_rng(0)
    acts = [gw.random_policy((2, 0), rng) for _ in range(4000)]
    assert all(isinstance(a, (int, np.integer)) for a in acts)
    counts = np.bincount(acts, minlength=4)
    assert counts.shape == (4,) and counts.min() > 900, f"all 4 actions roughly equally often, got {counts}"


def _non_terminal_states(env):
    return [(r, c) for r in range(env.n_rows) for c in range(env.n_cols) if env.grid[r][c] in ".S"]


def test_safe_policy_path():
    env = gw.GridWorld(slip=0.0, seed=0)
    for s in _non_terminal_states(env):
        a = gw.safe_policy(s, None)
        assert a in (0, 1, 2, 3), f"safe_policy must return an action in every state; got {a!r} at {s}"
    s = env.reset()
    visited = [s]
    for t in range(30):
        s, r, done = env.step(gw.safe_policy(s, None))
        visited.append(s)
        if done:
            break
    assert done and s == (2, 6), f"safe_policy should reach the goal; path was {visited}"
    assert len(visited) - 1 == 10, f"the long way round takes exactly 10 steps, took {len(visited) - 1}: {visited}"
    assert not any(r == 2 and 1 <= c <= 5 for r, c in visited), f"safe_policy should never walk on the bridge: {visited}"


# ---------------------------------------------------------------- Task 4
def test_run_episode_bridge_no_slip():
    env = gw.GridWorld(slip=0.0, seed=0)
    G, rewards = gw.run_episode(env, gw.bridge_policy, gamma=0.9)
    assert rewards == [0.0] * 5 + [1.0], f"bridge path: five 0s then +1, got {rewards}"
    assert G == pytest.approx(0.9**5), f"G_0 should be 0.9^5 = 0.59049, got {G}"


def test_run_episode_respects_max_steps():
    env = gw.GridWorld(slip=0.0, step_reward=-1.0, seed=0)
    always_left = lambda s, rng: LEFT  # bumps the edge forever
    G, rewards = gw.run_episode(env, always_left, gamma=1.0, max_steps=17)
    assert len(rewards) == 17, f"episode must stop after max_steps=17 steps, got {len(rewards)}"
    assert G == pytest.approx(-17.0)


def test_run_episode_passes_rng_to_policy():
    env = gw.GridWorld(slip=0.0, seed=0)
    G1, r1 = gw.run_episode(env, gw.random_policy, gamma=0.9, rng=np.random.default_rng(3))
    G2, r2 = gw.run_episode(env, gw.random_policy, gamma=0.9, rng=np.random.default_rng(3))
    assert r1 == r2 and G1 == G2, "same policy rng (and deterministic env) should give the same episode"


# ---------------------------------------------------------------- Task 5
def test_estimate_return_deterministic_case():
    env = gw.GridWorld(slip=0.0, seed=0)
    mean, se = gw.estimate_return(env, gw.bridge_policy, gamma=0.9, n_episodes=50)
    assert mean == pytest.approx(0.9**5), "deterministic env + policy: every return is 0.9^5"
    assert se == pytest.approx(0.0, abs=1e-12), "no spread, so the standard error is 0"


def test_estimate_return_matches_exact_values():
    # Exact values (computed by solving the Bellman equation, Chapter 3) at gamma = 0.9:
    #   slip 0.2: bridge -0.0274, safe 0.2637;  random policy (any slip): -0.4416
    env = gw.GridWorld(slip=0.2, seed=1)
    mb, sb = gw.estimate_return(env, gw.bridge_policy, gamma=0.9, n_episodes=3000, seed=2)
    ms, ss = gw.estimate_return(env, gw.safe_policy, gamma=0.9, n_episodes=3000, seed=3)
    assert abs(mb - (-0.0274)) < 4 * sb + 0.01, f"bridge at slip 0.2: expected about -0.027, got {mb:.3f} ± {sb:.3f}"
    assert abs(ms - 0.2637) < 4 * ss + 0.01, f"safe at slip 0.2: expected about 0.264, got {ms:.3f} ± {ss:.3f}"
    assert ms > mb, "with slip 0.2 the long way round should beat the bridge"
    assert 0 < sb < 0.05 and 0 < ss < 0.05, "standard errors should be small but positive"
    mr, sr = gw.estimate_return(gw.GridWorld(seed=4), gw.random_policy, gamma=0.9, n_episodes=3000, seed=5)
    assert abs(mr - (-0.4416)) < 4 * sr + 0.01, f"random policy: expected about -0.442, got {mr:.3f} ± {sr:.3f}"