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 _make_episode(T: int, state_dim: int, action_dim: int, value: float) -> Episode:
    states = torch.full((T, state_dim), value)
    actions = torch.zeros(T, action_dim)
    rewards = torch.full((T,), value)
    dones = torch.zeros(T, dtype=torch.bool)
    return Episode(states=states, actions=actions, rewards=rewards, dones=dones)


def test_add_episodes_batches_and_writes_once():
    state_dim, action_dim = 3, 2
    embed_dim = state_dim + action_dim + 1
    encoder = IdentityDualEncoder(input_dim=embed_dim)
    memory = MemoryStore(embed_dim=embed_dim, capacity=16)
    buffer = HippocampalReplayBuffer(memory=memory, encoder=encoder)

    episodes = [
        _make_episode(4, state_dim, action_dim, 1.0),
        _make_episode(6, state_dim, action_dim, 0.5),
        _make_episode(5, state_dim, action_dim, -0.2),
    ]

    buffer.add_episodes(episodes)

    assert len(memory) == len(episodes)
    assert memory.keys.shape == (len(episodes), embed_dim)


def test_add_episodes_updates_encoder_step_and_timestamps():
    state_dim, action_dim = 2, 1
    embed_dim = state_dim + action_dim + 1
    encoder = IdentityDualEncoder(input_dim=embed_dim)
    memory = MemoryStore(embed_dim=embed_dim, capacity=10)
    buffer = HippocampalReplayBuffer(memory=memory, encoder=encoder)

    episodes = [
        _make_episode(3, state_dim, action_dim, 0.1),
        _make_episode(4, state_dim, action_dim, 0.2),
        _make_episode(2, state_dim, action_dim, 0.3),
    ]
    buffer.add_episodes(episodes)

    # Encoder step should advance by number of ingested episodes
    step = int(buffer.stats()["encoder_step"])  # stored as float in stats()
    assert step == len(episodes)
    # All timestamps should equal the current step value
    assert memory.key_timestamps == [step] * len(episodes)
