import pytest

torch = pytest.importorskip("torch")

from hippotorch.core.episode import Episode


def _build_episode_rewards(rewards: list[float], dones: list[bool] | None = None) -> Episode:
    T = len(rewards)
    states = torch.randn(T, 3)
    actions = torch.randn(T, 2)
    r = torch.tensor(rewards, dtype=torch.float32)
    d = torch.tensor(dones, dtype=torch.bool) if dones is not None else None
    return Episode(states=states, actions=actions, rewards=r, dones=d)


def _mc_returns(rewards: list[float], gamma: float) -> list[float]:
    out = []
    for t in range(len(rewards)):
        s = 0.0
        g = 1.0
        for k in range(t, len(rewards)):
            s += g * rewards[k]
            g *= gamma
        out.append(s)
    return out


def test_n_step_returns_basic_and_with_dones() -> None:
    gamma = 0.9
    rewards = [1.0, 2.0, 3.0, 4.0]
    ep = _build_episode_rewards(rewards)

    # Monte-Carlo returns (no bootstrapping)
    ret_mc = ep.compute_returns(gamma, n_steps=None)
    expected_mc = torch.tensor(_mc_returns(rewards, gamma))
    assert torch.allclose(ret_mc, expected_mc, atol=1e-6)

    # 2-step returns without a value function -> zero bootstrap
    ret_2 = ep.compute_returns(gamma, n_steps=2)
    expected_2 = torch.tensor([1.0 + 0.9 * 2.0, 2.0 + 0.9 * 3.0, 3.0 + 0.9 * 4.0, 4.0])
    assert torch.allclose(ret_2, expected_2, atol=1e-6)

    # With a done at index 1, sums should truncate at the terminal
    ep_done = _build_episode_rewards(rewards, dones=[False, True, False, False])
    ret_done = ep_done.compute_returns(gamma, n_steps=2)
    expected_done = torch.tensor([1.0 + 0.9 * 2.0, 2.0, 3.0 + 0.9 * 4.0, 4.0])
    assert torch.allclose(ret_done, expected_done, atol=1e-6)


def test_gae_reduces_to_mc_when_lambda_one_and_zero_values() -> None:
    gamma = 0.95
    rewards = [0.5, -1.0, 2.0, 1.0]
    ep = _build_episode_rewards(rewards)

    # value function returning zero -> GAE with lambda=1 equals MC returns
    def zero_v(states: torch.Tensor) -> torch.Tensor:
        return torch.zeros(states.shape[0], dtype=torch.float32, device=states.device)

    ret_gae = ep.compute_returns(gamma, value_fn=zero_v, gae_lambda=1.0)
    expected_mc = torch.tensor(_mc_returns(rewards, gamma))
    assert torch.allclose(ret_gae, expected_mc, atol=1e-6)
    # advantages should be cached and equal to returns given zero baseline
    assert ep.advantages is not None
    assert torch.allclose(ep.advantages, ret_gae, atol=1e-6)
