from __future__ import annotations

import pytest

torch = pytest.importorskip("torch")

from hippotorch.core.episode import Episode
from hippotorch.memory.regression import IdentityDualEncoder
from hippotorch.memory.replay import HippocampalReplayBuffer
from hippotorch.memory.store import MemoryStore


def test_sample_with_intrinsic_attaches_intrinsic_field():
    state_dim, action_dim = 3, 2
    input_dim = state_dim + action_dim + 1
    encoder = IdentityDualEncoder(input_dim=input_dim)
    memory = MemoryStore(embed_dim=input_dim, capacity=8)
    buffer = HippocampalReplayBuffer(memory=memory, encoder=encoder, mixture_ratio=0.0)

    # Add a couple of short episodes
    for base in [0.0, 1.0]:
        T = 4
        states = torch.randn(T, state_dim) + base
        actions = torch.zeros(T, action_dim)
        rewards = torch.zeros(T)
        dones = torch.zeros(T, dtype=torch.bool)
        buffer.add_episode(Episode(states=states, actions=actions, rewards=rewards, dones=dones))

    batch = buffer.sample_with_intrinsic(batch_size=6, query_state=None, top_k=2)
    assert "intrinsic" in batch
    assert batch["intrinsic"].shape == batch["rewards"].shape
    # Intrinsic should be finite and non-negative in this simple setup
    assert torch.isfinite(batch["intrinsic"]).all()
    assert (batch["intrinsic"] >= 0).all()
