import pytest

torch = pytest.importorskip("torch")

from hippotorch.analysis.projector import MemoryProjector
from hippotorch.core.episode import Episode
from hippotorch.memory.regression import IdentityDualEncoder
from hippotorch.memory.store import MemoryStore


def _store_with_eps(n: int = 4, length: int = 4) -> MemoryStore:
    embed_dim = 4
    enc = IdentityDualEncoder(input_dim=embed_dim)
    mem = MemoryStore(embed_dim=enc.embed_dim, capacity=16)
    for i in range(n):
        v = float(i)
        states = torch.full((length, 2), v)
        actions = torch.zeros(length, 1)
        rewards = torch.ones(length) * v
        dones = torch.zeros(length, dtype=torch.bool)
        ep = Episode(states=states, actions=actions, rewards=rewards, dones=dones)
        tensor = mem.episodes_to_tensor([ep])
        keys = enc.encode_key(tensor)
        mem.write(keys, [ep], encoder_step=0)
    return mem


def test_memory_projector_pca_flow():
    mem = _store_with_eps(n=5)
    proj = MemoryProjector(sample_size=10, dims=2, method="pca")
    result = proj.project(mem)
    assert result.coords.shape[1] == 2
    assert len(result.contexts) == len(result.rewards) == result.coords.shape[0]


def test_memory_projector_with_context_fn():
    mem = _store_with_eps(n=3)
    proj = MemoryProjector(sample_size=10, dims=2, method="pca", context_fn=lambda ep: 7)
    result = proj.project(mem)
    assert all(c == 7 for c in result.contexts)


def test_memory_projector_invalid_method_raises():
    mem = _store_with_eps(n=3)
    proj = MemoryProjector(method="nope")
    with pytest.raises(ValueError):
        _ = proj.project(mem)


def test_memory_projector_empty_memory_returns_zeros():
    from hippotorch.memory.store import MemoryStore

    mem = MemoryStore(embed_dim=4, capacity=4)
    proj = MemoryProjector(sample_size=10, dims=2)
    result = proj.project(mem)
    assert result.coords.numel() == 0
    assert result.rewards == []
    assert result.contexts == []
