"""Tests for the Chapter 6 lab. Run one task at a time with  -k task1  etc.

The small cases are the same hand-computable numbers used in the chapter, so you
can check them with pen and paper.
"""

import copy
import importlib
import os

import numpy as np
import pytest
import torch
import torch.nn as nn

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

torch.set_num_threads(1)


class TableQ(nn.Module):
    """Q(s, .) = obs @ W. With one-hot observations this is a plain Q-table W[s, a],
    which makes losses and gradients easy to compute by hand."""

    def __init__(self, table):
        super().__init__()
        self.W = nn.Parameter(torch.tensor(table, dtype=torch.float32))

    def forward(self, obs):
        return obs @ self.W


def onehot(idx, n):
    return torch.eye(n, dtype=torch.float32)[torch.as_tensor(idx)]


# ------------------------------------------------------------------ Task 1
def test_task1_terminal_update_matches_hand_computation():
    # Chapter's worked example: "hat" features phi(x) = (1 - x, x) at x = 0.8,
    # target 1, alpha = 0.5, starting from w = 0.
    w = np.zeros(2)
    w_new = m.linear_td0_update(w, np.array([0.2, 0.8]), 1.0, np.zeros(2), True, 0.5, 0.9)
    np.testing.assert_allclose(w_new, [0.1, 0.4], atol=1e-12,
                               err_msg="delta = 1 - 0 = 1, so w = 0 + 0.5 * 1 * (0.2, 0.8)")
    assert np.all(w == 0), "linear_td0_update must not modify w in place"


def test_task1_bootstrapped_update_matches_hand_computation():
    # From x = 0.5 to x' = 0.8 with r = 0: V(s) = 0.25, V(s') = 0.34,
    # delta = 0.9 * 0.34 - 0.25 = 0.056, step = 0.5 * 0.056 * (0.5, 0.5).
    w = np.array([0.1, 0.4])
    w_new = m.linear_td0_update(w, np.array([0.5, 0.5]), 0.0, np.array([0.2, 0.8]), False, 0.5, 0.9)
    np.testing.assert_allclose(w_new, [0.114, 0.414], atol=1e-12,
                               err_msg="expected w + 0.5 * 0.056 * phi(s) = (0.114, 0.414)")


def test_task1_done_ignores_next_state():
    w = np.array([1.0, 2.0])
    a = m.linear_td0_update(w, np.array([1.0, 0.0]), 0.5, np.array([0.0, 1.0]), True, 0.1, 0.9)
    b = m.linear_td0_update(w, np.array([1.0, 0.0]), 0.5, np.array([7.0, -3.0]), True, 0.1, 0.9)
    np.testing.assert_allclose(a, b, err_msg="when done=True, V(s') must be treated as 0")


def test_task1_one_hot_features_give_tabular_td0():
    # With one-hot features, w IS the value table and the update is Chapter 5's TD(0).
    V = np.array([0.0, 0.5, 1.0])
    e = np.eye(3)
    V_new = m.linear_td0_update(V, e[0], 0.2, e[1], False, 0.1, 0.9)
    expected = V.copy()
    expected[0] += 0.1 * (0.2 + 0.9 * 0.5 - 0.0)
    np.testing.assert_allclose(V_new, expected, err_msg="one-hot features should reproduce tabular TD(0)")


def test_task1_deadly_triad_two_state_example():
    # V(A) = w, V(B) = 2w. Off-policy: only ever update on the A -> B transition (r = 0).
    # Each update multiplies w by 1 + alpha * (2 gamma - 1) = 1.08 for alpha = 0.1, gamma = 0.9.
    w = np.array([1.0])
    for _ in range(100):
        w = m.linear_td0_update(w, np.array([1.0]), 0.0, np.array([2.0]), False, 0.1, 0.9)
    np.testing.assert_allclose(w[0], 1.08 ** 100, rtol=1e-9,
                               err_msg="off-policy updates on A only should blow up like 1.08^100 ~ 2200")
    # On-policy: also update on B -> terminal (r = 0), as often as A. Now it converges.
    w = np.array([1.0])
    for _ in range(100):
        w = m.linear_td0_update(w, np.array([1.0]), 0.0, np.array([2.0]), False, 0.1, 0.9)
        w = m.linear_td0_update(w, np.array([2.0]), 0.0, np.array([0.0]), True, 0.1, 0.9)
    assert abs(w[0]) < 1e-3, f"with on-policy updates w should shrink to 0, got {w[0]}"


# ------------------------------------------------------------------ Task 2
def test_task2_push_and_len():
    buf = m.ReplayBuffer(capacity=5, obs_dim=2, seed=0)
    assert len(buf) == 0
    for i in range(3):
        buf.push(np.array([i, i]), i % 2, float(i), np.array([i + 1, i + 1]), i == 2)
    assert len(buf) == 3, "len should count stored transitions"
    assert buf.idx == 3, "idx should point at the next free slot"
    np.testing.assert_allclose(buf.obs[:3], [[0, 0], [1, 1], [2, 2]])
    np.testing.assert_allclose(buf.dones[:3], [0.0, 0.0, 1.0], err_msg="done must be stored as 0.0 / 1.0")


def test_task2_overwrites_oldest_when_full():
    buf = m.ReplayBuffer(capacity=4, obs_dim=1, seed=0)
    for i in range(6):  # transitions 0..5; 0 and 1 must be overwritten
        buf.push(np.array([i]), 0, float(i), np.array([i]), False)
    assert len(buf) == 4, "size must stop growing at capacity"
    assert buf.idx == 2, "after 6 pushes into capacity 4, the next slot is 6 % 4 = 2"
    assert sorted(buf.rewards.tolist()) == [2.0, 3.0, 4.0, 5.0], \
        "the two oldest transitions (rewards 0 and 1) should have been overwritten"


def test_task2_sample_shapes_and_dtypes():
    buf = m.ReplayBuffer(capacity=100, obs_dim=4, seed=0)
    for i in range(10):
        buf.push(np.full(4, i, np.float32), i % 2, 1.0, np.full(4, i + 1, np.float32), False)
    b = buf.sample(32)
    assert set(b) == {"obs", "actions", "rewards", "next_obs", "dones"}
    assert b["obs"].shape == (32, 4) and b["obs"].dtype == torch.float32
    assert b["next_obs"].shape == (32, 4) and b["next_obs"].dtype == torch.float32
    assert b["actions"].shape == (32,) and b["actions"].dtype == torch.int64
    assert b["rewards"].shape == (32,) and b["rewards"].dtype == torch.float32
    assert b["dones"].shape == (32,) and b["dones"].dtype == torch.float32
    # consistency: each sampled row is one stored transition
    assert torch.all(b["next_obs"][:, 0] == b["obs"][:, 0] + 1), "fields of a sample must come from the same slot"


def test_task2_sample_only_stored_and_roughly_uniform():
    buf = m.ReplayBuffer(capacity=1000, obs_dim=1, seed=1)
    for i in range(10):
        buf.push(np.array([i]), 0, float(i), np.array([i]), False)
    r = buf.sample(5000)["rewards"].numpy()
    assert r.min() >= 0 and r.max() <= 9, "sample must only return filled slots (indices < size)"
    counts = np.bincount(r.astype(int), minlength=10)
    assert counts.min() > 350 and counts.max() < 650, f"sampling should be ~uniform, got counts {counts}"


# ------------------------------------------------------------------ Task 3
def test_task3_qnetwork_architecture():
    torch.manual_seed(0)
    net = m.QNetwork(4, 2, hidden=64)
    n_params = sum(p.numel() for p in net.parameters())
    assert n_params == 4 * 64 + 64 + 64 * 64 + 64 + 64 * 2 + 2, \
        f"expected two hidden layers of 64 (4610 parameters), got {n_params}"
    out = net(torch.zeros(7, 4))
    assert out.shape == (7, 2), "forward should map (B, obs_dim) -> (B, n_actions)"
    n_linear = sum(isinstance(x, nn.Linear) for x in net.modules())
    assert n_linear == 3, "expected exactly 3 nn.Linear layers"


def test_task3_epsilon_schedule():
    f = m.epsilon_by_step
    assert f(0, 1.0, 0.05, 10_000) == pytest.approx(1.0)
    assert f(5_000, 1.0, 0.05, 10_000) == pytest.approx(0.525), "halfway: (1.0 + 0.05) / 2"
    assert f(10_000, 1.0, 0.05, 10_000) == pytest.approx(0.05)
    assert f(50_000, 1.0, 0.05, 10_000) == pytest.approx(0.05), "after decay_steps, hold at eps_end"


def test_task3_select_action_greedy_and_random():
    net = TableQ([[0.0, 1.0, -1.0], [2.0, 0.0, 0.0]])  # state 0 -> best action 1; state 1 -> 0
    rng = np.random.default_rng(0)
    s0, s1 = np.array([1.0, 0.0]), np.array([0.0, 1.0])
    a = m.select_action(net, s0, 0.0, rng)
    assert isinstance(a, int), "select_action must return a Python int"
    assert a == 1 and m.select_action(net, s1, 0.0, rng) == 0, "epsilon = 0 must be greedy"
    acts = [m.select_action(net, s0, 1.0, rng) for _ in range(3000)]
    counts = np.bincount(acts, minlength=3)
    assert counts.min() > 850, f"epsilon = 1 must be uniform over all 3 actions, got {counts}"


# ------------------------------------------------------------------ Task 4
# A 3-state, 2-action example (the one in the chapter's exercise):
#   transitions (s, a, r, s', done): (0,1, 1, 1,0), (1,0, 0, 2,0), (2,1, -1, 1,1); gamma = 0.5
TARGET_TABLE = [[0.0, 0.0], [1.0, 3.0], [2.0, -1.0]]
ONLINE_TABLE = [[0.0, 2.2], [5.0, 1.0], [0.0, 4.0]]


def make_batch():
    return {
        "obs": onehot([0, 1, 2], 3),
        "actions": torch.tensor([1, 0, 1]),
        "rewards": torch.tensor([1.0, 0.0, -1.0]),
        "next_obs": onehot([1, 2, 1], 3),
        "dones": torch.tensor([0.0, 0.0, 1.0]),
    }


def test_task4_targets_standard():
    b = make_batch()
    y = m.compute_td_targets(TableQ(TARGET_TABLE), b["rewards"], b["next_obs"], b["dones"], 0.5)
    torch.testing.assert_close(y, torch.tensor([2.5, 1.0, -1.0]),
                               msg="y = r + gamma (1-done) max_a' Q_target(s', a') = (1 + 0.5*3, 0 + 0.5*2, -1)")
    assert not y.requires_grad, "targets must be computed under torch.no_grad()"


def test_task4_targets_double():
    b = make_batch()
    y = m.compute_td_targets(TableQ(TARGET_TABLE), b["rewards"], b["next_obs"], b["dones"], 0.5,
                             online_net=TableQ(ONLINE_TABLE))
    # online picks a*=0 in s'=1 (5 > 1) and a*=1 in s'=2 (4 > 0); the target net evaluates them: 1 and -1
    torch.testing.assert_close(y, torch.tensor([1.5, -0.5, -1.0]),
                               msg="Double DQN: select a* with the online net, evaluate with the target net")
    assert not y.requires_grad


def test_task4_loss_value():
    b = make_batch()
    loss = m.dqn_loss(TableQ(ONLINE_TABLE), TableQ(TARGET_TABLE), b, 0.5)
    # Q(s,a) = (2.2, 5, 4); y = (2.5, 1, -1); errors (-0.3, 4, 5); Huber (0.045, 3.5, 4.5)
    assert loss.item() == pytest.approx(8.045 / 3, abs=1e-5), "mean Huber loss with threshold 1"
    loss_d = m.dqn_loss(TableQ(ONLINE_TABLE), TableQ(TARGET_TABLE), b, 0.5, double=True)
    # y = (1.5, -0.5, -1); errors (0.7, 5.5, 5); Huber (0.245, 5.0, 4.5)
    assert loss_d.item() == pytest.approx(9.745 / 3, abs=1e-5), "Double DQN loss"


def test_task4_gradient_is_semi_gradient():
    # Use the SAME network as online and target: the gradient must still flow only
    # through Q(s, a), never through the target (that's the "semi" in semi-gradient).
    b = make_batch()
    net = TableQ(ONLINE_TABLE)
    loss = m.dqn_loss(net, net, b, 0.5)
    loss.backward()
    # y = (1 + 0.5*5, 0 + 0.5*4, -1) = (3.5, 2, -1); errors (-1.3, 3, 5);
    # Huber'(err) = clip(err, -1, 1) = (-1, 1, 1); divided by B = 3.
    expected = torch.zeros(3, 2)
    expected[0, 1], expected[1, 0], expected[2, 1] = -1 / 3, 1 / 3, 1 / 3
    torch.testing.assert_close(net.W.grad, expected,
                               msg="gradient should be Huber'(Q - y) * dQ/dtheta / B, with y held constant")


def test_task4_no_gradient_into_target_net():
    b = make_batch()
    q, t = TableQ(ONLINE_TABLE), TableQ(TARGET_TABLE)
    m.dqn_loss(q, t, b, 0.5, double=True).backward()
    assert t.W.grad is None or torch.all(t.W.grad == 0), "no gradient may reach the target network"
    assert q.W.grad is not None and q.W.grad.abs().sum() > 0


# ------------------------------------------------------------------ Task 5
def test_task5_update_target_hard_and_soft():
    torch.manual_seed(0)
    q, t = m.QNetwork(4, 2, 16), m.QNetwork(4, 2, 16)
    m.update_target(t, q, tau=1.0)
    for p_t, p in zip(t.parameters(), q.parameters()):
        torch.testing.assert_close(p_t, p, msg="tau = 1 must copy the online weights")
    with torch.no_grad():
        for p in q.parameters():
            p.add_(1.0)
    for p_t, p in zip(t.parameters(), q.parameters()):
        assert not torch.allclose(p_t, p), "target must hold a COPY, not share storage with the online net"
    before = [p.detach().clone() for p in t.parameters()]
    m.update_target(t, q, tau=0.25)
    for p_t, p0, p in zip(t.parameters(), before, q.parameters()):
        torch.testing.assert_close(p_t, 0.25 * p + 0.75 * p0, msg="soft update: tau*theta + (1-tau)*theta_minus")


def test_task5_train_step_reduces_loss_and_leaves_target_alone():
    torch.manual_seed(0)
    q = m.QNetwork(4, 2, 32)
    t = copy.deepcopy(q)
    buf = m.ReplayBuffer(256, 4, seed=0)
    rng = np.random.default_rng(0)
    for _ in range(256):
        s = rng.normal(size=4).astype(np.float32)
        buf.push(s, int(rng.integers(2)), float(s[0] > 0), rng.normal(size=4).astype(np.float32), True)
    opt = torch.optim.Adam(q.parameters(), lr=1e-2)
    t_before = [p.detach().clone() for p in t.parameters()]
    losses = [m.train_step(q, t, opt, buf, 64, 0.99) for _ in range(150)]
    assert isinstance(losses[0], float), "train_step should return a Python float"
    assert np.mean(losses[-20:]) < 0.5 * np.mean(losses[:20]), \
        f"loss should fall on a fixed buffer with terminal transitions; went {np.mean(losses[:20]):.3f} -> {np.mean(losses[-20:]):.3f}"
    for p, p0 in zip(t.parameters(), t_before):
        torch.testing.assert_close(p, p0, msg="train_step must not change the target network")


# ------------------------------------------------------------------ Does it learn?
def test_task5_short_dqn_run_learns_cartpole():
    from train import Config, evaluate, train_dqn

    cfg = Config(total_steps=12_000, seed=0)
    out = train_dqn(m, cfg, eval_every=10**9)  # no intermediate evals
    score = evaluate(m, out["q_net"], n_episodes=10, seed=7)
    assert score > 60, (f"after 12k steps a working DQN balances the pole much longer than random "
                        f"(about 22 steps); got greedy return {score:.1f}")