"""Tests for the Chapter 18 lab. Run from this folder:

    python -m pytest -q                  # tests your agentic.py
    python -m pytest -q -k task2         # just one task
    RL_SOLUTIONS=1 python -m pytest -q   # tests the reference solution
"""

import importlib
import os

import numpy as np
import pytest
import torch

from tinyagent import (
    ANS, BOS, CALL, EQUAL, HIDDEN, HIGHER, LOWER, PAD, TOOL, V,
    Policy, make_batch, render,
)

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


# ----------------------------------------------------------------- Task 1: the environment
def test_task1_reset_shows_or_hides_the_secret():
    assert ag.GuessEnv(5, revealed=True).reset() == [BOS, 5]
    assert ag.GuessEnv(5, revealed=False).reset() == [BOS, HIDDEN]


def test_task1_tool_calls_compare_with_the_secret():
    env = ag.GuessEnv(5, revealed=False, call_cost=0.05)
    env.reset()
    assert env.step([CALL, 7]) == ([TOOL, LOWER], -0.05, False), "secret 5 < 7: the tool should say 'lower'"
    assert env.step([CALL, 2]) == ([TOOL, HIGHER], -0.05, False), "secret 5 > 2: the tool should say 'higher'"
    assert env.step([CALL, 5]) == ([TOOL, EQUAL], -0.05, False), "secret 5 == 5: the tool should say 'equal'"


def test_task1_answers_end_the_episode():
    env = ag.GuessEnv(3, revealed=False)
    env.reset()
    assert env.step([ANS, 3]) == ([], 1.0, True), "a correct answer earns 1.0 and ends the episode"
    env = ag.GuessEnv(3, revealed=False)
    env.reset()
    assert env.step([ANS, 4]) == ([], 0.0, True), "a wrong answer earns 0.0 and ends the episode"


@pytest.mark.parametrize("turn", [[CALL, CALL], [ANS, TOOL], [3, 4], [LOWER, 2], [CALL, HIDDEN]])
def test_task1_malformed_turns_are_penalised_and_halt(turn):
    env = ag.GuessEnv(3, revealed=False, malformed_reward=-1.0)
    env.reset()
    obs, r, done = env.step(turn)
    assert obs == [] and done, f"malformed turn {turn} should halt the episode with no tool output"
    assert r == -1.0, f"malformed turn {turn} should earn malformed_reward = -1.0, got {r}"


def test_task1_running_out_of_turns():
    env = ag.GuessEnv(6, revealed=False, max_turns=3)
    env.reset()
    assert env.step([CALL, 1])[2] is False
    assert env.step([CALL, 2])[2] is False
    obs, r, done = env.step([CALL, 3])
    assert done, "the 3rd turn of a max_turns=3 episode must end it, even if it was a tool call"
    assert obs == [TOOL, HIGHER], "the last tool call still returns its output"


# ----------------------------------------------------------------- Task 2: the rollout loop
class ScriptedSampler:
    """A fake sampler that plays a fixed list of model tokens and records what it was fed."""

    ENV_TOKENS = {TOOL, LOWER, HIGHER, EQUAL}

    def __init__(self, plan):
        self.plan, self.fed = plan, []

    def start(self):
        return ()

    def step(self, state, tok):
        self.fed.append(tok)
        state = state + (tok,)
        n_model = sum(1 for t in state[2:] if t not in self.ENV_TOKENS)   # tokens after the prompt
        nxt = self.plan[n_model] if n_model < len(self.plan) else ANS
        probs = np.zeros(V)
        probs[nxt] = 1.0
        return probs, state


def test_task2_rollout_records_who_wrote_each_token():
    # secret 4, hidden; the "policy" plays a binary search: call 3, call 5, answer 4
    sampler = ScriptedSampler([CALL, 3, CALL, 5, ANS, 4])
    env = ag.GuessEnv(4, revealed=False, call_cost=0.05)
    tr = ag.rollout(sampler, env, np.random.default_rng(0))
    assert tr.tokens == [BOS, HIDDEN, CALL, 3, TOOL, HIGHER, CALL, 5, TOOL, LOWER, ANS, 4], \
        f"wrong token sequence: {render(tr)}"
    assert tr.is_model == [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1], \
        "is_model must be 1 exactly for the tokens the policy sampled (not prompt, not tool output)"
    assert tr.turn_id == [-1, -1, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2], \
        "turn_id: -1 for the prompt, k for turn k's model tokens AND the tool output that answers them"
    assert tr.turn_rewards == pytest.approx([-0.05, -0.05, 1.0]), "one reward per turn"


def test_task2_rollout_feeds_every_token_to_the_policy():
    sampler = ScriptedSampler([CALL, 3, CALL, 5, ANS, 4])
    tr = ag.rollout(sampler, ag.GuessEnv(4, revealed=False), np.random.default_rng(0))
    assert sampler.fed == tr.tokens, (
        "the sampler must be fed every token in order (prompt, model tokens AND tool output): "
        "the policy has to read what the tool said")


def test_task2_rollout_with_a_real_policy_is_consistent():
    torch.manual_seed(0)
    sampler = Policy().sampler()
    rng = np.random.default_rng(1)
    for k in range(50):
        tr = ag.rollout(sampler, ag.GuessEnv(k % 8, revealed=bool(k % 2)), rng)
        n = len(tr.tokens)
        assert len(tr.is_model) == n and len(tr.turn_id) == n
        assert sum(tr.is_model) == 2 * tr.n_turns, "each turn has exactly two model tokens"
        n_tool_outputs = sum(1 for i in range(n) if tr.tokens[i] == TOOL and not tr.is_model[i])
        assert n == 2 + 2 * tr.n_turns + 2 * n_tool_outputs, "each tool call adds two tool tokens"
        assert tr.n_turns <= 5


# ----------------------------------------------------------------- Task 3: advantages
GROUP = [[-0.05, -0.05, 1.0],   # two calls, then right:   R = 0.90
         [1.0],                 # answered right at once:  R = 1.00
         [-0.05, 0.0],          # one call, then wrong:    R = -0.05
         [-1.0]]                # malformed:               R = -1.00


def test_task3_trajectory_level_advantages_by_hand():
    adv = ag.turn_advantages(GROUP, mode="trajectory")
    b = (0.9 + 1.0 - 0.05 - 1.0) / 4                    # 0.2125
    expected = [[0.9 - b] * 3, [1.0 - b], [-0.05 - b] * 2, [-1.0 - b]]
    assert len(adv) == 4
    for a, e in zip(adv, expected):
        np.testing.assert_allclose(a, e, atol=1e-12,
                                   err_msg="trajectory mode: every turn gets R_i - mean(R)")


def test_task3_turn_level_advantages_by_hand():
    adv = ag.turn_advantages(GROUP, mode="turn", gamma=1.0)
    b = 0.2125
    expected = [[0.9 - b, 0.95 - b, 1.0 - b], [1.0 - b], [-0.05 - b, 0.0 - b], [-1.0 - b]]
    for a, e in zip(adv, expected):
        np.testing.assert_allclose(a, e, atol=1e-12,
                                   err_msg="turn mode: turn k gets (reward-to-go from k) - mean(G_0)")


def test_task3_turn_level_with_discount():
    adv = ag.turn_advantages([[0.0, 0.0, 1.0], [1.0]], mode="turn", gamma=0.5)
    # returns-to-go: [0.25, 0.5, 1.0] and [1.0]; baseline = mean(0.25, 1.0) = 0.625
    np.testing.assert_allclose(adv[0], [0.25 - 0.625, 0.5 - 0.625, 1.0 - 0.625], atol=1e-12)
    np.testing.assert_allclose(adv[1], [1.0 - 0.625], atol=1e-12)


def test_task3_zero_signal_group():
    for mode in ("trajectory", "turn"):
        adv = ag.turn_advantages([[1.0], [1.0], [1.0]], mode=mode)
        for a in adv:
            np.testing.assert_allclose(a, [0.0], atol=1e-12,
                                       err_msg="if every trajectory earns the same, nobody gets pushed")


# ----------------------------------------------------------------- Task 4: the masked loss
def _tiny_batch():
    torch.manual_seed(0)
    ids = torch.tensor([[BOS, HIDDEN, CALL, 3, TOOL, HIGHER, ANS, 4],
                        [BOS, 2, ANS, 2, PAD, PAD, PAD, PAD]])
    is_model = torch.tensor([[0, 0, 1, 1, 0, 0, 1, 1],
                             [0, 0, 1, 1, 0, 0, 0, 0]], dtype=torch.float)
    adv = torch.tensor([[0.0, 0.0, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5],
                        [0.0, 0.0, -1.0, -1.0, 0.0, 0.0, 0.0, 0.0]])
    logits = torch.randn(2, 7, V, requires_grad=True)
    return ids, is_model, adv, logits


def test_task4_loss_value_matches_the_formula():
    ids, mask, adv, logits = _tiny_batch()
    loss = ag.masked_policy_loss(logits, ids, mask, adv)
    logp = torch.log_softmax(logits, -1)
    terms = []
    for b in range(2):
        for t in range(1, 8):
            if mask[b, t] > 0:
                terms.append(-adv[b, t] * logp[b, t - 1, ids[b, t]])
    expected = torch.stack(terms).sum() / len(terms)          # 6 model tokens in total
    assert torch.isfinite(loss), "loss must be finite (padding must not leak in)"
    assert loss.item() == pytest.approx(expected.item(), abs=1e-6), \
        "loss = -sum(mask * A * log pi(next token)) / sum(mask), with logits[:, t-1] predicting ids[:, t]"


def test_task4_tool_tokens_get_exactly_zero_gradient():
    ids, mask, adv, logits = _tiny_batch()
    ag.masked_policy_loss(logits, ids, mask, adv).backward()
    g = logits.grad
    # logits[:, t-1] predicts ids[:, t]; tool tokens of row 0 are at t = 4, 5 -> logit rows 3, 4
    assert torch.all(g[0, 3] == 0) and torch.all(g[0, 4] == 0), \
        "positions that predict tool-output tokens must receive no gradient"
    assert torch.all(g[1, 3:] == 0), "padding must receive no gradient"
    assert g[0, 1].abs().sum() > 0 and g[0, 5].abs().sum() > 0, "model tokens must receive gradient"


def test_task4_without_the_mask_tool_tokens_do_get_gradient():
    ids, mask, adv, logits = _tiny_batch()
    everything = (ids != PAD).float()
    everything[:, :2] = 0                                   # still exclude the prompt
    ag.masked_policy_loss(logits, ids, everything, adv).backward()
    assert logits.grad[0, 3].abs().sum() > 0, "with a mask that includes tool tokens, they get gradient"


def test_task4_empty_mask_is_safe():
    ids, mask, adv, logits = _tiny_batch()
    loss = ag.masked_policy_loss(logits, ids, torch.zeros_like(mask), adv)
    assert torch.isfinite(loss) and loss.item() == 0.0, "an all-zero mask should give a loss of 0"


def test_task4_make_batch_and_loss_work_together():
    torch.manual_seed(0)
    pol = Policy()
    rng = np.random.default_rng(0)
    trajs = [ag.rollout(pol.sampler(), ag.GuessEnv(k % 8, False), rng) for k in range(8)]
    advs = ag.turn_advantages([t.turn_rewards for t in trajs])
    ids, mask, adv = make_batch(trajs, advs)
    loss = ag.masked_policy_loss(pol(ids[:, :-1]), ids, mask, adv)
    assert torch.isfinite(loss)
    loss.backward()


# ----------------------------------------------------------------- does it learn?
def test_training_learns_to_use_the_tool():
    _, hist = ag.train(iters=100, tasks_per_iter=24, warm_steps=500, seed=0)
    first = np.mean([h["solved_hard"] for h in hist[:10]])
    last = np.mean([h["solved_hard"] for h in hist[-10:]])
    calls = np.mean([h["calls_hard"] for h in hist[-10:]])
    assert last > first + 0.15, (
        f"on hidden-number tasks the success rate should climb (first 10 iters {first:.2f}, "
        f"last 10 {last:.2f})")
    assert calls > 1.0, f"the agent should be calling the tool on hidden-number tasks (got {calls:.2f} calls)"