from __future__ import annotations

import torch

from pi06mini.experience import (
    build_synthetic_experience,
    discounted_returns,
    n_step_advantages,
    normalize_returns_to_unit_interval,
    quantile_advantage_labels,
)


def _geometry() -> tuple[torch.Tensor, torch.Tensor]:
    history = torch.stack((torch.linspace(-7.0, 0.0, 8), torch.zeros(8)), dim=-1)
    x = torch.linspace(1.0, 16.0, 16)
    future = torch.stack((x, 0.025 * x.square()), dim=-1)
    return history, future


def test_synthetic_buffer_has_all_experience_sources_and_failure_geometry() -> None:
    history, future = _geometry()
    episodes = build_synthetic_experience(history, future, episode_prefix="pai")

    assert [episode.source_type for episode in episodes] == [
        "demo_success",
        "autonomous_understeer",
        "autonomous_stop",
        "intervention_correction",
    ]
    demo, understeer, stopped, correction = episodes
    torch.testing.assert_close(demo.future_xy, future)
    assert understeer.future_xy[-1, 1] < future[-1, 1]
    assert torch.equal(stopped.future_xy[-1], stopped.future_xy[-2])
    assert correction.intervention_mask.any()
    assert not correction.intervention_mask[0]
    torch.testing.assert_close(correction.future_xy[-1], future[-1])
    assert torch.all(demo.rewards == 0)
    assert float(understeer.rewards.min()) < 0
    assert float(stopped.rewards.min()) < 0


def test_returns_are_normalized_to_closed_negative_unit_interval() -> None:
    rewards = torch.tensor([[0.0, 0.0, 0.0], [-0.2, -0.4, -1.0]])
    returns = discounted_returns(rewards, gamma=1.0)
    normalized = normalize_returns_to_unit_interval(returns)

    assert normalized.shape == rewards.shape
    assert float(normalized.min()) == -1.0
    assert float(normalized.max()) == 0.0
    assert torch.all(normalized <= 0.0)
    assert torch.all(normalized >= -1.0)
    assert torch.equal(
        normalize_returns_to_unit_interval(torch.ones(3)), torch.zeros(3)
    )


def test_n_step_advantage_bootstraps_and_stops_at_terminal() -> None:
    rewards = torch.tensor([1.0, 2.0, 3.0])
    values = torch.tensor([0.5, 1.0, 1.5, 8.0])
    dones = torch.tensor([False, True, False])
    advantages = n_step_advantages(
        rewards, values, n_step=3, gamma=0.5, dones=dones
    )

    # t=0 sees r0 + gamma*r1, then the terminal blocks the large bootstrap.
    assert torch.isclose(advantages[0], torch.tensor(1.5))
    assert torch.isclose(advantages[1], torch.tensor(1.0))
    # t=2 reaches the explicit end and therefore has no bootstrap.
    assert torch.isclose(advantages[2], torch.tensor(1.5))


def test_quantile_labels_force_interventions_positive() -> None:
    advantages = torch.tensor([-3.0, -2.0, -1.0, 2.0])
    interventions = torch.tensor([True, False, False, False])
    labels, threshold = quantile_advantage_labels(
        advantages, quantile=0.75, intervention_mask=interventions
    )

    assert labels.tolist() == [True, False, False, True]
    assert float(threshold) > -1.0