import pytest

torch = pytest.importorskip("torch")

from hippotorch import (Episode, HippocampalReplayBuffer, IdentityDualEncoder,
                        MemoryStore, query)


def _make_episode(value: float = 0.0) -> Episode:
    states = torch.full((4, 2), value)
    actions = torch.zeros(4, 1)
    rewards = torch.full((4,), value + 1.0)
    dones = torch.zeros(4, dtype=torch.bool)
    return Episode(states=states, actions=actions, rewards=rewards, dones=dones)


def test_query_is_read_only():
    embed_dim = 4  # 2 state dims + 1 action + 1 reward
    encoder = IdentityDualEncoder(input_dim=embed_dim)
    memory = MemoryStore(embed_dim=embed_dim, capacity=8)
    buffer = HippocampalReplayBuffer(memory=memory, encoder=encoder, mixture_ratio=0.0)
    for offset in [0.0, 1.0]:
        buffer.add_episode(_make_episode(value=offset))

    baseline_access = list(memory.access_counts)
    baseline_timestamps = list(memory.key_timestamps)
    baseline_versions = list(memory.slot_versions)

    obs = torch.tensor([0.5, 0.5, 0.0, 0.0], dtype=torch.float32)
    result = query(obs, buffer=buffer, top_k=2, return_scores=False)

    assert len(result.episodes[0]) <= 2
    assert result.scores is None
    assert memory.access_counts == baseline_access
    assert memory.key_timestamps == baseline_timestamps
    assert memory.slot_versions == baseline_versions


def test_query_accepts_memory_encoder_pair():
    embed_dim = 4
    encoder = IdentityDualEncoder(input_dim=embed_dim)
    memory = MemoryStore(embed_dim=embed_dim, capacity=4)
    buffer = HippocampalReplayBuffer(memory=memory, encoder=encoder, mixture_ratio=0.0)
    buffer.add_episode(_make_episode(value=2.0))

    obs = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float32)
    result = query(
        obs,
        memory=memory,
        encoder=encoder,
        top_k=1,
        return_scores=True,
        return_attention=True,
    )
    assert result.scores is not None
    assert result.attention is not None
    assert len(result.episode_ids[0]) == 1
