from __future__ import annotations

import torch

from hippotorch.memory.regression import IdentityDualEncoder
from hippotorch.memory.replay import HippocampalReplayBuffer
from hippotorch.memory.sb3 import SB3ReplayBufferWrapper
from hippotorch.memory.store import MemoryStore


def test_sb3_sampling_diagnostics_reports_semantic_when_query_active():
    # Identity-like encoder ensures semantic proximity is meaningful
    input_dim = 3 + 1 + 1  # obs(3) + act(1) + reward(1)
    encoder = IdentityDualEncoder(input_dim=input_dim)
    memory = MemoryStore(embed_dim=input_dim, capacity=8)
    buffer = HippocampalReplayBuffer(memory=memory, encoder=encoder, mixture_ratio=1.0)
    wrapper = SB3ReplayBufferWrapper(buffer)

    # Two short episodes with different observation scales
    o0 = torch.full((3,), 1.0)
    o1 = torch.full((3,), 2.0)
    a = torch.zeros(1)
    wrapper.add(o0, o1, a, reward=1.0, done=False)
    wrapper.add(o1, o1 + 1.0, a, reward=1.0, done=True)

    o0b = torch.full((3,), -1.0)
    o1b = torch.full((3,), -2.0)
    wrapper.add(o0b, o1b, a, reward=0.1, done=False)
    wrapper.add(o1b, o1b - 1.0, a, reward=0.1, done=True)

    # Query near the positive cluster should trigger semantic sampling
    q = torch.full((3,), 2.0)
    _ = wrapper.sample(16, query_observation=q, top_k=1)
    stats = wrapper.sampling_diagnostics()
    assert stats.get("semantic", 0.0) > 0.0
