from __future__ import annotations

import torch

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


def test_next_state_alignment_in_sample_batch():
    # Build buffer with uniform sampling only
    state_dim, action_dim = 3, 1
    input_dim = state_dim + action_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=0.0)

    # Create one episode with strictly increasing states so alignment is obvious
    T = 6
    states = torch.stack([torch.full((state_dim,), float(i)) for i in range(T)], dim=0)
    actions = torch.zeros(T, action_dim)
    rewards = torch.arange(T, dtype=torch.float)
    dones = torch.zeros(T, dtype=torch.bool)

    from hippotorch.core.episode import Episode

    buffer.add_episode(Episode(states=states, actions=actions, rewards=rewards, dones=dones))

    batch = buffer.sample(batch_size=16)
    s, ns = batch["states"], batch["next_states"]
    # For every sampled transition, next state must be exactly +1 step forward
    # in the original sequence since our episode has no terminal steps and
    # sampling never chooses the final step.
    # Decode the scalar marker from the first coordinate
    diffs = (ns[:, 0] - s[:, 0]).cpu().numpy()
    assert (diffs == 1.0).all()
