import torch

from pi0mini.model import Pi0Config, Pi0Mini


def test_pi0_three_block_attention_mask() -> None:
    config = Pi0Config(
        image_size=32,
        patch_size=16,
        max_text_tokens=4,
        max_action_tokens=3,
        d_model=32,
        n_heads=4,
        n_layers=1,
        time_feature_dim=8,
    )
    model = Pi0Mini(config)

    # Two image patches per axis -> 4 image tokens, plus 2 language tokens.
    prefix_edges, suffix_edges = model._masks(
        batch=1,
        prefix_length=6,
        suffix_length=4,  # one state + three action tokens
        image_tokens=4,
        attention_mask=torch.tensor([[1, 1]]),
        device=torch.device("cpu"),
    )

    assert prefix_edges.shape == (1, 6, 6)
    assert torch.all(prefix_edges)

    # Suffix key order is [six prefix keys, state, action0, action1, action2].
    assert suffix_edges.shape == (1, 4, 10)
    assert torch.all(suffix_edges[0, 0, :7])
    assert not torch.any(suffix_edges[0, 0, 7:])
    assert torch.all(suffix_edges[0, 1:, :])


def test_pi0_mask_respects_padded_language_keys() -> None:
    model = Pi0Mini(
        Pi0Config(
            image_size=16,
            patch_size=16,
            max_text_tokens=3,
            max_action_tokens=1,
            d_model=16,
            n_heads=4,
            n_layers=1,
            time_feature_dim=8,
        )
    )
    prefix_edges, suffix_edges = model._masks(
        batch=1,
        prefix_length=4,
        suffix_length=2,
        image_tokens=1,
        attention_mask=torch.tensor([[1, 0, 0]]),
        device=torch.device("cpu"),
    )

    # Padded prefix keys are hidden from valid prefix, state, and action rows.
    assert not prefix_edges[0, 0, 2]
    assert not prefix_edges[0, 0, 3]
    assert not suffix_edges[0, :, 2].any()
    assert not suffix_edges[0, :, 3].any()