from __future__ import annotations

import pytest
import torch

from pi07mini.rtc import align_actions_for_rtc, sample_rtc_delay


def _actions() -> torch.Tensor:
    return torch.arange(32, dtype=torch.float32).reshape(16, 2)


def test_zero_delay_preserves_entire_chunk() -> None:
    actions = _actions()
    aligned = align_actions_for_rtc(actions, 0)
    torch.testing.assert_close(aligned.delayed_target_actions, actions)
    assert aligned.valid_mask.all()
    assert aligned.executed_prefix_actions.shape == (0, 2)


def test_delay_shifts_target_and_masks_padded_tail() -> None:
    actions = _actions()
    aligned = align_actions_for_rtc(actions, 12)
    torch.testing.assert_close(aligned.delayed_target_actions[:4], actions[12:])
    assert torch.count_nonzero(aligned.delayed_target_actions[4:]) == 0
    assert aligned.valid_mask.tolist() == [True] * 4 + [False] * 12
    torch.testing.assert_close(aligned.executed_prefix_actions, actions[:12])


def test_rtc_alignment_preserves_batch_dimensions() -> None:
    actions = torch.stack((_actions(), _actions() + 100.0))
    aligned = align_actions_for_rtc(actions, 3)
    assert aligned.delayed_target_actions.shape == (2, 16, 2)
    assert aligned.valid_mask.shape == (2, 16)
    torch.testing.assert_close(aligned.delayed_target_actions[:, :13], actions[:, 3:])


def test_delay_sampling_covers_closed_zero_to_twelve_range() -> None:
    generator = torch.Generator().manual_seed(9)
    values = {
        sample_rtc_delay(generator=generator, action_horizon=16) for _ in range(300)
    }
    assert min(values) == 0
    assert max(values) == 12


@pytest.mark.parametrize("delay", [-1, 13, 16])
def test_invalid_delay_is_rejected(delay: int) -> None:
    with pytest.raises(ValueError):
        align_actions_for_rtc(_actions(), delay)