from __future__ import annotations

import pytest
import torch

from pi07mini.flow import FlowMatching
from pi07mini.rtc import build_rtc_conditioning, sample_rtc_delay


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


def test_zero_delay_has_empty_prefix_and_full_postfix() -> None:
    rtc = build_rtc_conditioning(_actions(), 0)
    assert not rtc.prefix_mask.any()
    assert rtc.postfix_mask.all()
    assert torch.count_nonzero(rtc.action_prefix) == 0


def test_delay_keeps_ground_truth_prefix_without_shifting_target() -> None:
    actions = _actions()
    rtc = build_rtc_conditioning(actions, 12)
    torch.testing.assert_close(rtc.action_prefix[:12], actions[:12])
    assert torch.count_nonzero(rtc.action_prefix[12:]) == 0
    assert rtc.prefix_mask.tolist() == [True] * 12 + [False] * 4
    assert rtc.postfix_mask.tolist() == [False] * 12 + [True] * 4


def test_flow_training_sets_prefix_time_to_one_and_masks_its_loss() -> None:
    actions = _actions()[None]
    noisy, times, _, loss_mask = FlowMatching().training_batch(
        actions, delays=torch.tensor([3]), generator=torch.Generator().manual_seed(4)
    )
    torch.testing.assert_close(noisy[:, :3], actions[:, :3])
    torch.testing.assert_close(times[:, :3], torch.ones(1, 3))
    assert not loss_mask[:, :3].any()
    assert loss_mask[:, 3:].all()


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):
        build_rtc_conditioning(_actions(), delay)