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_transition_alignment_next_state_matches():
    obs_dim = 3
    act_dim = 1
    input_dim = obs_dim + act_dim + 1

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

    # Build two episodes where states[t+1] = states[t] + 1 (per-dim)
    for base in [0.0, 10.0]:
        for t in range(3):
            obs_t = torch.full((obs_dim,), base + t)
            obs_tp1 = torch.full((obs_dim,), base + t + 1)
            action = torch.zeros(act_dim)
            done = t == 2
            wrapper.add(obs_t, obs_tp1, action, reward=0.0, done=done)

    # Sample many transitions uniformly; verify next_states == states + 1
    batch = wrapper.sample(batch_size=16)
    states = batch["states"]
    next_states = batch["next_states"]
    dones = batch["dones"]

    # No sampled transition should be terminal because sampling avoids the last step
    assert not bool(dones.any().item())
    assert torch.allclose(next_states, states + 1.0)
