from __future__ import annotations

import torch

from hippotorch.agents.sac import SACAgent, SACConfig
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 _make_continuous_episode(T: int, state_dim: int, action_dim: int) -> Episode:
    states = torch.randn(T, state_dim)
    actions = torch.randn(T, action_dim).clamp(-1.0, 1.0)
    rewards = torch.randn(T)
    dones = torch.zeros(T, dtype=torch.bool)
    dones[-1] = True
    return Episode(states=states, actions=actions, rewards=rewards, dones=dones)


def test_sac_agent_accepts_replay_batch() -> None:
    torch.manual_seed(0)
    state_dim, action_dim = 7, 3
    # Encoder expects obs + action + reward
    encoder = DualEncoder(input_dim=state_dim + action_dim + 1, embed_dim=32)
    memory = MemoryStore(embed_dim=32, capacity=1000)
    buffer = HippocampalReplayBuffer(memory=memory, encoder=encoder, mixture_ratio=0.0)

    episodes = [_make_continuous_episode(8, state_dim, action_dim) for _ in range(20)]
    buffer.add_episodes(episodes)

    batch = buffer.sample(batch_size=32)

    agent = SACAgent(SACConfig(state_dim=state_dim, action_dim=action_dim, action_limit=1.0))
    metrics = agent.update(batch)

    assert "pi_loss" in metrics and "q1_loss" in metrics
    assert isinstance(metrics["pi_loss"], float)


def test_sac_agent_act_shape_and_bounds() -> None:
    torch.manual_seed(0)
    state_dim, action_dim = 5, 2
    agent = SACAgent(SACConfig(state_dim=state_dim, action_dim=action_dim, action_limit=2.0))
    state = torch.randn(state_dim)
    action = agent.act(state, explore=True)
    assert action.shape == (action_dim,)
    assert (action <= 2.0 + 1e-6).all() and (action >= -2.0 - 1e-6).all()
