import pytest

torch = pytest.importorskip("torch")

from hippotorch.core.episode import Episode
from hippotorch.memory.regression import IdentityDualEncoder
from hippotorch.memory.store import MemoryStore
from hippotorch.utils.diagnostics import log_memory_pca, pca_keys, pca_keys_with_rewards


def test_pca_shapes_and_wrappers():
    keys = torch.randn(16, 8)
    comps, coords = pca_keys(keys, k=2)
    assert comps.shape == (8, 2)
    assert coords.shape == (16, 2)

    rewards = torch.linspace(0, 1, steps=16)
    payload = pca_keys_with_rewards(keys, rewards, k=2)
    assert set(payload.keys()) == {"components", "coords", "rewards"}
    assert payload["coords"].shape[1] == 2


def test_log_memory_pca_invokes_writer():
    # Build a tiny memory with two episodes
    embed_dim = 4
    enc = IdentityDualEncoder(input_dim=embed_dim)
    mem = MemoryStore(embed_dim=enc.embed_dim, capacity=8)
    for v in [0.0, 1.0]:
        states = torch.full((4, 2), v)
        actions = torch.zeros(4, 1)
        rewards = torch.ones(4) * v
        dones = torch.zeros(4, 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)

    called = {"n": 0}

    class FakeWriter:
        def add_embedding(
            self, coords, metadata=None, tag=None, global_step=None
        ):  # noqa: D401 - test stub
            called["n"] += 1

    log_memory_pca(FakeWriter(), mem, step=0, sample_size=2)
    assert called["n"] == 1
