from __future__ import annotations

import pytest
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 _make_episode(length: int, obs_val: float, act_dim: int = 2):
    states = torch.full((length, 3), obs_val)
    actions = torch.zeros(length, act_dim)
    rewards = torch.full((length,), 1.0 if obs_val > 0 else 0.1)
    dones = torch.zeros(length, dtype=torch.bool)
    return states, actions, rewards, dones


@pytest.mark.parametrize("mixture", [1.0])
def test_sb3_semantic_activation_improves_reward(mixture: float):
    # Deterministic encoder: Identity-like mapping ensures semantic proximity.
    input_dim = 3 + 2 + 1
    encoder = IdentityDualEncoder(input_dim=input_dim)
    memory = MemoryStore(embed_dim=input_dim)
    buf_sem = HippocampalReplayBuffer(memory=memory, encoder=encoder, mixture_ratio=mixture)
    wrapper = SB3ReplayBufferWrapper(buf_sem)

    # Build a small dataset: good episodes around +1, distractors around -1.
    # Record two 2-step episodes of each kind via wrapper.add.
    obs0 = torch.full((3,), 1.0)
    obs1 = torch.full((3,), 2.0)
    a = torch.zeros(2)
    wrapper.add(obs0, obs1, a, reward=1.0, done=False)
    wrapper.add(obs1, obs1 + 1.0, a, reward=1.0, done=True)

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

    # Uniform buffer for comparison
    mem2 = MemoryStore(embed_dim=input_dim)
    buf_uni = HippocampalReplayBuffer(memory=mem2, encoder=encoder, mixture_ratio=0.0)
    w2 = SB3ReplayBufferWrapper(buf_uni)
    # Re-add same episodes
    w2.add(obs0, obs1, a, reward=1.0, done=False)
    w2.add(obs1, obs1 + 1.0, a, reward=1.0, done=True)
    w2.add(obs0b, obs1b, a, reward=0.1, done=False)
    w2.add(obs1b, obs1b - 1.0, a, reward=0.1, done=True)

    # Query near the + side should retrieve higher-reward transitions.
    q = torch.full((3,), 2.0)
    sem_batch = wrapper.sample(32, query_observation=q, top_k=1)
    uni_batch = w2.sample(32)
    sem_mean = float(sem_batch["rewards"].mean().item())
    uni_mean = float(uni_batch["rewards"].mean().item())
    assert sem_mean >= uni_mean - 1e-8
