from __future__ import annotations

import torch

from hippotorch.encoders.dual import DualEncoder
from hippotorch.memory.replay import HippocampalReplayBuffer
from hippotorch.memory.sb3 import SB3ReplayBufferWrapper
from hippotorch.memory.store import MemoryStore


def test_sb3_wrapper_uniform_fallback_when_no_query():
    obs_dim = 2
    act_dim = 1
    input_dim = obs_dim + act_dim + 1

    encoder = DualEncoder(input_dim=input_dim, embed_dim=8, momentum=0.9, refresh_interval=10)
    memory = MemoryStore(embed_dim=8, capacity=10)
    buffer = HippocampalReplayBuffer(memory=memory, encoder=encoder, mixture_ratio=1.0)
    wrapper = SB3ReplayBufferWrapper(buffer)

    o0 = torch.zeros(obs_dim)
    o1 = torch.ones(obs_dim)
    a = torch.tensor([0.5])

    wrapper.add(o0, o1, a, reward=1.0, done=False)
    wrapper.add(o1, o1 + 1, a, reward=0.5, done=True)

    # After done, no last observation retained, so query is None.
    _batch = wrapper.sample(8)
    stats = wrapper.sampling_diagnostics()
    # With mixture=1.0 and no query, semantic count should be 0 (full uniform fallback)
    assert stats.get("semantic", 0.0) == 0.0
