import torch

from alpamayo_lite.trajectory_tokens import (
    DiscreteTrajectoryMask,
    TrajectoryVocabulary,
    dequantize_action,
    extract_trajectory_values,
    quantize_action,
)


def test_quantized_action_round_trip_is_bounded() -> None:
    action = torch.tensor([[-4.0, -1.25], [0.0, 1.5], [3.9, 4.0]])
    tokens = quantize_action(action, size=128, minimum=-4.0, maximum=4.0)
    reconstructed = dequantize_action(tokens, size=128, minimum=-4.0, maximum=4.0)
    assert float((reconstructed - action).abs().max()) <= 8.0 / 127.0


def test_extract_trajectory_values_between_markers() -> None:
    vocabulary = TrajectoryVocabulary(
        size=8,
        start_id=100,
        end_id=107,
        special_token_ids={"<|traj_future_start|>": 200, "<|traj_future_end|>": 201},
    )
    sequence = torch.tensor([1, 200, 103, 107, 100, 201, 2])
    torch.testing.assert_close(
        extract_trajectory_values(sequence, vocabulary), torch.tensor([3, 7, 0])
    )


def test_discrete_trajectory_mask_changes_only_trajectory_range() -> None:
    input_ids = torch.tensor([[1, 2]])
    scores = torch.randn(1, 16)
    original = scores.clone()
    masked = DiscreteTrajectoryMask(start_id=5, size=4)(input_ids, scores)
    assert torch.isneginf(masked[:, 5:9]).all()
    torch.testing.assert_close(masked[:, :5], original[:, :5])
    torch.testing.assert_close(masked[:, 9:], original[:, 9:])