from __future__ import annotations

import pytest

torch = pytest.importorskip("torch")

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


def test_buffer_stats_keys_stable_across_operations():
    state_dim, action_dim = 3, 1
    input_dim = state_dim + action_dim + 1
    encoder = DualEncoder(input_dim=input_dim, embed_dim=8, momentum=0.0, refresh_interval=10)
    memory = MemoryStore(embed_dim=8, capacity=8)
    buffer = HippocampalReplayBuffer(memory=memory, encoder=encoder, mixture_ratio=0.5)

    # Add a couple of small episodes
    for _ in range(2):
        T = 4
        ep = Episode(
            states=torch.randn(T, state_dim),
            actions=torch.randn(T, action_dim),
            rewards=torch.randn(T),
            dones=torch.zeros(T, dtype=torch.bool),
        )
        buffer.add_episode(ep, encoder_step=0)

    # Baseline stats and required keys
    s0 = buffer.stats()
    required = {"memory_size", "semantic_fraction", "encoder_step", "sampled_steps"}
    assert required.issubset(s0.keys())

    # Perform typical operations: sample and (if consolidator absent) just resample
    q = torch.cat([torch.randn(state_dim), torch.zeros(action_dim), torch.zeros(1)])
    try:
        _ = buffer.sample(batch_size=3, query_state=q, top_k=2)
    except Exception:
        pass

    s1 = buffer.stats()
    # Keys should be stable (may gain extra fields like staleness_*, but not lose baseline)
    assert required.issubset(s1.keys())
    assert set(s0.keys()).issubset(set(s1.keys())) or set(s1.keys()).issubset(set(s0.keys()))
