import torch

from pi07mini.tokenizer import EpisodeMetadata, MiniTokenizer


def metadata() -> EpisodeMetadata:
    return EpisodeMetadata(
        speed_steps=2100,
        quality=5,
        mistake=False,
        control_mode="ee",
    )


def test_prompt_contains_rich_metadata_and_control_tokens() -> None:
    tokenizer = MiniTokenizer()
    prompt, metadata_indices = tokenizer.encode_prompt(
        "drive on the road", "steer_left", metadata()
    )
    assert tokenizer.SPEED_MARKER in prompt
    assert tokenizer.QUALITY_MARKER in prompt
    assert tokenizer.MISTAKE_MARKER in prompt
    assert tokenizer.CONTROL_MARKER in prompt
    assert tokenizer.CONTROL_EE in prompt
    assert len(metadata_indices) == 6


def test_high_level_prompt_contains_past_subtasks_and_next_query() -> None:
    tokenizer = MiniTokenizer()
    prompt = tokenizer.encode_high_level_prompt(
        "drive", ("keep_lane", "slow_down"), metadata()
    )
    assert prompt[-1] == tokenizer.NEXT_SUBTASK_MARKER
    assert torch.count_nonzero(prompt == tokenizer.SUBTASK_MARKER) == 2


def test_metadata_cfg_pair_drops_metadata_but_retains_control_mode() -> None:
    tokenizer = MiniTokenizer()
    actions = torch.zeros(3, 2)
    conditioned, unconditional = tokenizer.build_metadata_cfg_pair(
        "drive", "keep_lane", actions, metadata()
    )
    assert tokenizer.QUALITY_MARKER in conditioned.input_ids
    assert tokenizer.QUALITY_MARKER not in unconditional.input_ids
    assert tokenizer.CONTROL_EE in conditioned.input_ids
    assert tokenizer.CONTROL_EE in unconditional.input_ids
    assert conditioned.metadata_indices
    assert not unconditional.metadata_indices


def test_flow_prefix_ends_before_fast_stream() -> None:
    tokenizer = MiniTokenizer()
    sequence = tokenizer.build_policy_sequence(
        "drive", "keep_lane", torch.zeros(4, 2), metadata()
    )
    assert sequence.input_ids[sequence.action_start_index] == tokenizer.ACTION_START
    assert sequence.action_prefix_mask[: sequence.action_start_index].all()
    assert not sequence.action_prefix_mask[sequence.action_start_index :].any()
    assert sequence.labels[sequence.action_start_index - 1] == 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