import torch

from pi07mini.memory import MEMConfig, MEMVideoEncoder


def tiny_encoder() -> MEMVideoEncoder:
    return MEMVideoEncoder(
        MEMConfig(
            image_size=16,
            patch_size=8,
            d_model=16,
            n_heads=4,
            n_layers=2,
            max_history_frames=6,
        )
    ).eval()


def test_mem_output_token_count_is_independent_of_history_length() -> None:
    encoder = tiny_encoder()
    current = torch.randn(2, 3, 16, 16)
    with torch.no_grad():
        no_history = encoder(current)
        long_history = encoder(current, torch.randn(2, 6, 3, 16, 16))
    assert no_history.shape == (2, 4, 16)
    assert long_history.shape == no_history.shape


def test_mem_is_sensitive_to_temporal_order() -> None:
    torch.manual_seed(7)
    encoder = tiny_encoder()
    current = torch.randn(1, 3, 16, 16)
    history = torch.stack(
        (
            torch.zeros(3, 16, 16),
            torch.ones(3, 16, 16),
            torch.full((3, 16, 16), 2.0),
        )
    )[None]
    with torch.no_grad():
        chronological = encoder(current, history)
        reversed_history = encoder(current, history.flip(1))
    assert not torch.allclose(chronological, reversed_history)


def test_masked_history_frame_cannot_change_memory() -> None:
    torch.manual_seed(9)
    encoder = tiny_encoder()
    current = torch.randn(1, 3, 16, 16)
    history = torch.randn(1, 2, 3, 16, 16)
    changed = history.clone()
    changed[:, 0] = 1000.0
    mask = torch.tensor([[False, True]])
    with torch.no_grad():
        first = encoder(current, history, mask)
        second = encoder(current, changed, mask)
    torch.testing.assert_close(first, second)


def test_current_frame_temporal_position_is_exactly_zero() -> None:
    encoding = MEMVideoEncoder.temporal_position(torch.tensor([3, 2, 1, 0]), 15)
    assert torch.count_nonzero(encoding[-1]) == 0