import torch

from pi06mini.tokenizer import MiniTokenizer


def test_advantage_condition_is_between_subtask_and_actions() -> None:
    tokenizer = MiniTokenizer()
    history = torch.tensor([[0.0, 0.0], [1.0, 0.2]])
    actions = torch.tensor([[0.1, 0.0], [0.2, -0.1]])

    positive = tokenizer.build_policy_sequence(
        history, "steer_left", actions, "positive"
    )
    assert positive.advantage_index == positive.subtask_index + 1
    assert positive.action_start_index == positive.advantage_index + 1
    assert positive.input_ids[positive.advantage_index] == tokenizer.ADVANTAGE_POSITIVE
    assert positive.input_ids[positive.action_start_index] == tokenizer.ACTION_START
    assert positive.action_prefix_mask[: positive.action_start_index].all()
    assert not positive.action_prefix_mask[positive.action_start_index :].any()

    negative = tokenizer.build_policy_sequence(
        history, "steer_left", actions, "negative"
    )
    assert negative.input_ids[negative.advantage_index] == tokenizer.ADVANTAGE_NEGATIVE


def test_omitted_advantage_is_absent_not_a_third_condition_token() -> None:
    tokenizer = MiniTokenizer()
    sequence = tokenizer.build_policy_sequence(
        torch.zeros(2, 2), "keep_lane", torch.zeros(2, 2), "omitted"
    )

    assert sequence.advantage_index is None
    assert sequence.action_start_index == sequence.subtask_index + 1
    assert sequence.input_ids[sequence.action_start_index] == tokenizer.ACTION_START
    assert tokenizer.ADVANTAGE_POSITIVE not in sequence.input_ids
    assert tokenizer.ADVANTAGE_NEGATIVE not in sequence.input_ids


def test_labels_predict_subtask_and_fast_but_not_external_advantage() -> None:
    tokenizer = MiniTokenizer()
    positive = tokenizer.build_policy_sequence(
        torch.zeros(2, 2), "stop", torch.zeros(2, 2), True
    )
    assert positive.labels[positive.subtask_index - 1] == positive.input_ids[positive.subtask_index]
    assert positive.labels[positive.subtask_index] == -100
    assert positive.labels[positive.advantage_index] == tokenizer.ACTION_START

    omitted = tokenizer.build_policy_sequence(
        torch.zeros(2, 2), "stop", torch.zeros(2, 2), None
    )
    assert omitted.labels[omitted.subtask_index] == tokenizer.ACTION_START


def test_fast_lite_dct_round_trip_has_bounded_error() -> None:
    tokenizer = MiniTokenizer()
    actions = torch.tensor(
        [[0.0, 0.0], [0.2, -0.1], [0.5, -0.2], [0.8, -0.1]],
        dtype=torch.float32,
    )
    exact = tokenizer.idct_actions(tokenizer.dct_actions(actions))
    torch.testing.assert_close(exact, actions, atol=1e-5, rtol=1e-5)

    decoded = tokenizer.decode_actions(tokenizer.encode_actions(actions))
    assert torch.sqrt(torch.mean((decoded - actions) ** 2)) < 0.3