import torch

from hippotorch.core.episode import Episode
from hippotorch.encoders.visual import VisualEpisodeEncoder
from hippotorch.memory.prototype import PrototypeExtractor
from hippotorch.memory.regression import IdentityDualEncoder
from hippotorch.memory.store import MemoryStore


def _make_episode(state_value: float, reward_value: float, *, length: int = 3) -> Episode:
    states = torch.full((length, 2), state_value)
    actions = torch.zeros(length, 1)
    rewards = torch.full((length,), reward_value)
    return Episode(states=states, actions=actions, rewards=rewards)


def test_prototype_extraction_compresses_memory() -> None:
    torch.manual_seed(0)
    episodes = [
        _make_episode(1.0, 1.0),
        _make_episode(1.1, 1.2),
        _make_episode(0.9, 0.9),
        _make_episode(-1.0, 0.1),
    ]

    encoder = IdentityDualEncoder(input_dim=3)
    memory = MemoryStore(embed_dim=encoder.projector.out_features)
    tensors = memory.episodes_to_tensor(episodes)
    memory.write(encoder.encode_key(tensors), episodes)

    extractor = PrototypeExtractor(kmeans_iters=10, keep_per_cluster=1)
    result = extractor.extract_prototypes(
        memory, encoder, num_prototypes=2, reward_percentile=0.5, prune_sources=True
    )

    # Expect three candidates above the median total reward
    assert result.num_candidates == 3
    # With k=2 and keep_per_cluster=1, we should retain two medoids
    assert result.num_prototypes == 2

    # Memory size should decrease by the number of pruned candidates
    assert len(memory) == len(episodes) - len(result.indices_pruned)

    # Verify medoids are truly nearest to their cluster centroids on candidate keys
    # Recompute candidate keys deterministically from the original episodes
    tensors = memory.episodes_to_tensor(episodes)
    keys = encoder.encode_key(tensors)
    rewards = torch.tensor([ep.total_reward for ep in episodes], dtype=torch.float)
    thr = torch.quantile(rewards, q=0.5, interpolation="lower")
    candidate_indices = [i for i, r in enumerate(rewards.tolist()) if r >= float(thr)]
    cand_keys = keys[candidate_indices]

    def _kmeans_full_batch(
        x: torch.Tensor, k: int, iters: int
    ) -> tuple[torch.Tensor, torch.Tensor]:
        N = x.shape[0]
        # Deterministic farthest-first init
        chosen = [0]
        while len(chosen) < k:
            current = x[torch.tensor(chosen)]
            d2 = ((x.unsqueeze(1) - current.unsqueeze(0)) ** 2).sum(dim=-1)
            min_d2 = d2.min(dim=1).values
            min_d2[torch.tensor(chosen)] = -1.0
            next_idx = int(torch.argmax(min_d2).item())
            if next_idx in chosen:
                break
            chosen.append(next_idx)
        centroids = x[torch.tensor(chosen[:k])].clone()
        assign = torch.zeros(N, dtype=torch.long)
        for _ in range(max(1, iters)):
            # squared distances
            x_norm = (x**2).sum(dim=1, keepdim=True)
            c_norm = (centroids**2).sum(dim=1, keepdim=True).t()
            d2 = x_norm + c_norm - 2 * (x @ centroids.t())
            assign = torch.argmin(d2, dim=1)
            new_centroids = centroids.clone()
            for j in range(k):
                m = assign == j
                if torch.any(m):
                    new_centroids[j] = x[m].mean(dim=0)
            centroids = new_centroids
        return centroids, assign

    centroids, assignments = _kmeans_full_batch(cand_keys, k=2, iters=10)
    expected_keep: list[int] = []
    for j in range(2):
        m = assignments == j
        if not torch.any(m):
            continue
        members = cand_keys[m]
        d = torch.cdist(members.unsqueeze(0), centroids[j : j + 1].unsqueeze(0)).squeeze()
        member_positions = torch.nonzero(m, as_tuple=False).squeeze(-1)
        # keep_per_cluster = 1
        if members.shape[0] == 1:
            best = member_positions[0].item()
        else:
            order = torch.argsort(d)
            best = member_positions[order[0]].item()
        expected_keep.append(candidate_indices[best])

    actual_kept = sorted(set(candidate_indices) - set(result.indices_pruned))
    assert set(expected_keep) == set(actual_kept)


def test_visual_encoder_nature_cnn_shape() -> None:
    encoder = VisualEpisodeEncoder(image_channels=3, embed_dim=16)
    single = torch.randint(0, 255, (4, 3, 48, 48), dtype=torch.uint8)
    outputs_single = encoder(single)
    assert outputs_single.shape == (1, single.shape[0], encoder.embed_dim)
    norms_single = outputs_single.norm(dim=-1)
    assert torch.allclose(norms_single, torch.ones_like(norms_single), atol=1e-5)

    batched = torch.rand(3, 5, 3, 48, 48)
    outputs_batch = encoder(batched)
    assert outputs_batch.shape == (3, 5, encoder.embed_dim)
    norms = outputs_batch.norm(dim=-1)
    assert torch.allclose(norms, torch.ones_like(norms), atol=1e-5)
