import torch

from pi05mini.model import Pi05Config, Pi05Mini


def tiny_model() -> Pi05Mini:
    return Pi05Mini(Pi05Config(
        vocab_size=300,
        image_size=16,
        patch_size=16,
        d_model=32,
        n_heads=4,
        n_layers=1,
        max_text_tokens=16,
        max_action_tokens=4,
    ))


def test_pi05_has_no_continuous_state_projection() -> None:
    model = tiny_model()
    assert not hasattr(model, "state_projection")
    output = model(
        torch.randn(2, 3, 16, 16),
        torch.randint(0, 300, (2, 7)),
        torch.randn(2, 4, 2),
        torch.rand(2),
        generated_token_start=4,
    )
    assert output.text_logits.shape == (2, 7, 300)
    assert output.vector_field.shape == (2, 4, 2)


def test_observation_generated_and_flow_attention_partition() -> None:
    model = tiny_model()
    # One image patch followed by four observation tokens and three generated
    # tokens. Flow may read observation plus the first generated (subtask)
    # token, but not the two teacher-forced FAST tokens.
    action_prefix_mask = torch.tensor([[1, 1, 1, 1, 1, 1, 0, 0]], dtype=torch.bool)
    prefix, action = model._masks(
        batch=1,
        prefix_length=8,
        action_length=2,
        text_offset=1,
        attention_mask=torch.ones(1, 7, dtype=torch.bool),
        action_prefix_mask=action_prefix_mask,
        generated_token_start=4,
        device=torch.device("cpu"),
    )

    generated_start = 5  # image patch + four observation text tokens
    # Observation block is fully bidirectional and cannot see generated tokens.
    assert prefix[0, 0, :generated_start].all()
    assert not prefix[0, 0, generated_start:].any()
    assert prefix[0, 4, :generated_start].all()
    # Generated tokens see all observation tokens and only prior/self generated tokens.
    assert prefix[0, 5, :6].all() and not prefix[0, 5, 6:].any()
    assert prefix[0, 6, :7].all() and not prefix[0, 6, 7]
    # Flow reads allowed prefix keys and all flow-action keys bidirectionally.
    assert action[0, :, :6].all()
    assert not action[0, :, 6:8].any()
    assert action[0, :, 8:].all()