from __future__ import annotations

import pytest
import torch

from pi05mini.tokenizer import MiniTokenizer


def test_orthonormal_dct_round_trip() -> None:
    tokenizer = MiniTokenizer()
    actions = torch.tensor(
        [
            [-4.0, -2.75],
            [-0.25, 0.0],
            [1.375, 3.9],
            [4.0, 2.125],
        ]
    )

    coefficients = tokenizer.dct_actions(actions)
    reconstructed = tokenizer.idct_actions(coefficients)

    assert coefficients.shape == actions.shape
    torch.testing.assert_close(reconstructed, actions, atol=1e-5, rtol=1e-5)


def test_fast_lite_tokens_round_trip_with_bounded_quantization_error() -> None:
    tokenizer = MiniTokenizer()
    actions = torch.tensor(
        [
            [-3.5, -2.75],
            [-0.25, 0.0],
            [1.375, 3.5],
            [3.75, 2.125],
        ]
    )

    tokens = tokenizer.encode_actions(actions)
    reconstructed = tokenizer.decode_actions(tokens)

    bin_width = (
        tokenizer.ACTION_MAXIMUM - tokenizer.ACTION_MINIMUM
    ) / (tokenizer.ACTION_BINS - 1)
    assert reconstructed.shape == actions.shape
    # Orthonormal IDCT preserves total coefficient error.  Coefficients are
    # quantized after division by sqrt(horizon), yielding this RMSE bound.
    rmse = torch.sqrt(torch.mean((reconstructed - actions).square()))
    assert float(rmse) <= bin_width / 2 * actions.shape[0] ** 0.5 + 1e-6


def test_action_tokens_are_serialized_low_frequency_first() -> None:
    tokenizer = MiniTokenizer()
    horizon = 4
    coefficients = torch.zeros(horizon, 2)
    coefficients[0, 0] = 1.5 * horizon**0.5  # positive X DC coefficient
    coefficients[1, 1] = -1.5 * horizon**0.5  # negative Y first harmonic
    actions = tokenizer.idct_actions(coefficients)

    payload = tokenizer.encode_actions(actions)[1:-1]
    zero_token = tokenizer.ACTION_TOKEN_START + round(
        (0.0 - tokenizer.ACTION_MINIMUM)
        / (tokenizer.ACTION_MAXIMUM - tokenizer.ACTION_MINIMUM)
        * (tokenizer.ACTION_BINS - 1)
    )

    # Frequency-major XY order is [DC-x, DC-y, f1-x, f1-y, ...].
    assert payload[0].item() > zero_token
    assert payload[1].item() == zero_token
    assert payload[2].item() == zero_token
    assert payload[3].item() < zero_token
    assert torch.all(payload[4:] == zero_token)


def test_action_token_ids_are_contiguous_reserved_vocabulary() -> None:
    tokenizer = MiniTokenizer()
    edge_actions = torch.tensor([[-4.0, 4.0], [0.0, 1.0]])
    tokens = tokenizer.encode_actions(edge_actions)

    assert tokens[0].item() == tokenizer.ACTION_START
    assert tokens[-1].item() == tokenizer.ACTION_END
    assert tokens[1:-1].min().item() >= tokenizer.ACTION_TOKEN_START
    assert tokens[1:-1].max().item() <= tokenizer.ACTION_TOKEN_END
    assert tuple(tokenizer.action_token_ids) == tuple(
        range(tokenizer.ACTION_TOKEN_START, tokenizer.ACTION_TOKEN_END + 1)
    )
    assert tokenizer.vocab_size == tokenizer.ACTION_TOKEN_END + 1


def test_decode_actions_rejects_missing_boundaries_and_non_action_ids() -> None:
    tokenizer = MiniTokenizer()

    with pytest.raises(ValueError, match="enclosed"):
        tokenizer.decode_actions(torch.tensor([tokenizer.ACTION_TOKEN_START] * 2))
    with pytest.raises(ValueError, match="reserved"):
        tokenizer.decode_actions(
            torch.tensor(
                [
                    tokenizer.ACTION_START,
                    tokenizer.BOS,
                    tokenizer.ACTION_TOKEN_START,
                    tokenizer.ACTION_END,
                ]
            )
        )