import pytest

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

from hippotorch.core.episode import Episode
from hippotorch.memory.store import MemoryStore
from hippotorch.utils.ann import (_FAISS_AVAILABLE,
                                  ApproximateNearestNeighborIndex, FaissIndex)


@pytest.mark.skipif(not _FAISS_AVAILABLE, reason="FAISS not installed")
def test_faiss_matches_exact_search():
    dim = 16
    keys = torch.randn(256, dim)
    queries = torch.randn(4, dim)

    faiss_index = FaissIndex()
    faiss_index.rebuild(keys)
    scores_faiss, idx_faiss = faiss_index.search(queries, top_k=5)

    torch_index = ApproximateNearestNeighborIndex(use_faiss=False)
    torch_index.rebuild(keys)
    scores_torch, idx_torch = torch_index.search(queries, top_k=5)

    assert torch.allclose(scores_faiss, scores_torch, atol=1e-5)
    assert torch.equal(idx_faiss, idx_torch)


def _make_episode(vec) -> Episode:
    state_vec = torch.as_tensor(vec, dtype=torch.float32)
    states = state_vec.unsqueeze(0).repeat(3, 1)
    actions = torch.zeros(3, 1)
    rewards = torch.zeros(3, dtype=torch.float32)
    dones = torch.zeros(3, dtype=torch.bool)
    return Episode(states=states, actions=actions, rewards=rewards, dones=dones)


class _ScalarEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.online = nn.Identity()
        self.target = nn.Identity()
        self.dummy = nn.Parameter(torch.zeros(1))

    def encode_key(self, x):
        return x.mean(dim=1)

    def prepare_episode(self, episode: Episode) -> torch.Tensor:
        return episode.states


@pytest.mark.parametrize("use_faiss", [False, True])
def test_memory_index_refresh_and_eviction(use_faiss: bool) -> None:
    if use_faiss and not _FAISS_AVAILABLE:
        pytest.skip("FAISS unavailable")
    indexer = ApproximateNearestNeighborIndex(use_faiss=use_faiss)
    memory = MemoryStore(embed_dim=2, capacity=8, indexer=indexer)
    episodes = [
        _make_episode((1.0, 0.0)),
        _make_episode((0.0, 1.0)),
        _make_episode((0.5, 0.5)),
    ]
    keys = torch.tensor([[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]], dtype=torch.float32)
    memory.write(keys, episodes)

    encoder = _ScalarEncoder()
    memory.episodes[1] = _make_episode((0.0, -1.0))
    memory.refresh_keys([1], encoder=encoder, encoder_step=42)

    _, idx = memory.indexer.search(torch.tensor([[0.0, -1.0]]), top_k=1)
    assert idx.item() == 1

    removed_key = memory.keys[0].clone()
    memory.prune_indices([0])
    _, idx_after = memory.indexer.search(torch.tensor([[1.0, 0.0]]), top_k=1)
    returned = memory.keys[idx_after.item()]
    assert not torch.allclose(returned, removed_key)


@pytest.mark.parametrize("use_faiss", [False, True])
def test_memory_rebuild_index_deterministic(use_faiss: bool) -> None:
    if use_faiss and not _FAISS_AVAILABLE:
        pytest.skip("FAISS unavailable")
    indexer = ApproximateNearestNeighborIndex(use_faiss=use_faiss)
    memory = MemoryStore(embed_dim=2, capacity=8, indexer=indexer)
    episodes = [
        _make_episode((1.0, 0.0)),
        _make_episode((0.0, 1.0)),
        _make_episode((0.5, 0.5)),
    ]
    keys = torch.tensor([[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]], dtype=torch.float32)
    memory.write(keys, episodes)

    target_idx = 2
    new_keys = memory.keys.clone()
    new_keys[target_idx] = torch.tensor([0.0, -1.0], dtype=torch.float32)
    memory.keys = new_keys
    _, stale_idx = memory.indexer.search(torch.tensor([[0.0, -1.0]]), top_k=1)
    assert stale_idx.item() != target_idx

    memory.rebuild_index()
    _, synced_idx = memory.indexer.search(torch.tensor([[0.0, -1.0]]), top_k=1)
    assert synced_idx.item() == target_idx
