import pytest

torch = pytest.importorskip("torch")
from torch import nn

from hippotorch.core.episode import Episode
from hippotorch.encoders import (
    DualEncoder,
    MemoryEncoder,
    VisualEpisodeEncoder,
    get_encoder,
    list_encoders,
    register_encoder,
)
from hippotorch.memory.replay import HippocampalReplayBuffer
from hippotorch.memory.store import MemoryStore


def test_builtin_encoders_satisfy_protocol():
    denc = DualEncoder(input_dim=8, embed_dim=4)
    venc = VisualEpisodeEncoder(image_channels=3, embed_dim=8)
    assert isinstance(denc, MemoryEncoder)
    assert isinstance(venc, MemoryEncoder)


class TrivialEncoder(nn.Module):
    """Minimal encoder satisfying the MemoryEncoder protocol for testing.

    Pools inputs over time and applies a linear projection for both query and
    key paths. Maintains a simple refresh counter to satisfy the lifecycle
    hooks used by the buffer/memory.
    """

    def __init__(self, input_dim: int, embed_dim: int, refresh_interval: int = 10) -> None:
        super().__init__()
        self.projector = nn.Linear(int(input_dim), int(embed_dim), bias=False)
        nn.init.eye_(self.projector.weight)  # type: ignore[arg-type]
        self.refresh_interval = int(refresh_interval)
        self._since = 0

    def encode_query(self, x, mask=None):  # type: ignore[override]
        pooled = x.mean(dim=1)
        return self.projector(pooled)

    def encode_key(self, x, mask=None):  # type: ignore[override]
        pooled = x.mean(dim=1)
        return self.projector(pooled)

    def update_target(self) -> None:  # type: ignore[override]
        self._since += 1

    def should_refresh_keys(self) -> bool:  # type: ignore[override]
        return self._since >= self.refresh_interval

    def mark_refreshed(self) -> None:  # type: ignore[override]
        self._since = 0


def _mk_episode(T=4, Ds=3, Da=2, r=1.0) -> Episode:
    s = torch.randn(T, Ds)
    a = torch.randn(T, Da)
    rw = torch.full((T,), float(r))
    d = torch.zeros(T, dtype=torch.bool)
    return Episode(states=s, actions=a, rewards=rw, dones=d)


def test_pluggable_registry_and_end_to_end_usage():
    name = "trivial"
    register_encoder(name, lambda input_dim, embed_dim: TrivialEncoder(input_dim, embed_dim))
    assert name in list_encoders()

    factory = get_encoder(name)
    Ds, Da = 3, 2
    embed_dim = Ds + Da + 1
    enc = factory(input_dim=embed_dim, embed_dim=embed_dim)
    assert isinstance(enc, MemoryEncoder)

    mem = MemoryStore(embed_dim=enc.projector.out_features, capacity=16)
    buf = HippocampalReplayBuffer(memory=mem, encoder=enc, mixture_ratio=1.0)

    episodes = [_mk_episode(T=6, Ds=Ds, Da=Da, r=1.0) for _ in range(6)]
    for i, ep in enumerate(episodes):
        buf.add_episode(ep, encoder_step=i)

    # Build a simple query from the first episode's initial transition
    ep0 = episodes[0]
    q = torch.cat([ep0.states[0], ep0.actions[0], ep0.rewards[0].unsqueeze(0)])
    batch = buf.sample(batch_size=4, query_state=q, top_k=3)
    assert set(batch.keys()) == {"states", "actions", "rewards", "next_states", "dones"}
    assert batch["states"].shape[0] > 0
