import pytest

torch = pytest.importorskip("torch")

from hippotorch.core.episode import Episode
from hippotorch.memory.regression import IdentityDualEncoder
from hippotorch.memory.replay import HippocampalReplayBuffer
from hippotorch.memory.store import MemoryStore


def _make_episode(T: int, state_dim: int, action_dim: int, value: float) -> Episode:
    states = torch.full((T, state_dim), value)
    actions = torch.zeros(T, action_dim)
    rewards = torch.full((T,), value)
    dones = torch.zeros(T, dtype=torch.bool)
    return Episode(states=states, actions=actions, rewards=rewards, dones=dones)


def test_priority_weighted_sampling_prefers_high_priority_transitions():
    state_dim, action_dim = 3, 2
    embed_dim = state_dim + action_dim + 1
    encoder = IdentityDualEncoder(input_dim=embed_dim)
    memory = MemoryStore(embed_dim=embed_dim, capacity=32)

    # Create 3 episodes with 4 steps each (3 transitions per episode)
    episodes = [_make_episode(4, state_dim, action_dim, 0.1 * (i + 1)) for i in range(3)]
    # Write keys via encoder
    for ep in episodes:
        x = torch.cat([ep.states, ep.actions, ep.rewards.unsqueeze(-1)], dim=-1).unsqueeze(0)
        key = encoder.encode_key(x)
        memory.write(key, [ep], encoder_step=0)

    # Mark all transitions low priority except those from episode 0
    hi_indices, hi_errors = [], []
    lo_indices, lo_errors = [], []
    for ei, ep in enumerate(episodes):
        for ti in range(len(ep) - 1):
            if ei == 0:  # high-priority episode
                hi_indices.append((ei, ti))
                hi_errors.append(10.0)
            else:
                lo_indices.append((ei, ti))
                lo_errors.append(0.1)

    memory.update_priorities(hi_indices, torch.tensor(hi_errors))
    memory.update_priorities(lo_indices, torch.tensor(lo_errors))

    # Build buffer with prioritized mode, but set mixture to 0 to avoid semantic branch
    buffer = HippocampalReplayBuffer(
        memory=memory,
        encoder=encoder,
        mixture_ratio=0.0,
        priority_mode="proportional",
        priority_alpha=1.0,
        priority_beta=0.0,
    )

    # Draw many samples and count how often each transition is seen
    counts = {(ei, ti): 0 for ei in range(3) for ti in range(3)}
    num_draws = 1000
    # Use store's sampler to recover indices alongside sampling
    for _ in range(num_draws):
        samples, _, idx = memory.sample_priority_transitions(
            1, mode="proportional", alpha=1.0, beta=0.0
        )
        if idx.numel() == 0:
            continue
        pair = (int(idx[0, 0].item()), int(idx[0, 1].item()))
        counts[pair] += 1

    hi_avg = sum(counts[(0, ti)] for ti in range(3)) / 3.0
    lo_avg = sum(counts[(ei, ti)] for ei in (1, 2) for ti in range(3)) / 6.0

    # High-priority transitions should be sampled substantially more often
    assert hi_avg > lo_avg * 2.0, f"expected high-priority > 2x low, got {hi_avg} vs {lo_avg}"

    # Buffer.sample should attach weights when priority mode is enabled
    batch = buffer.sample(batch_size=8, query_state=None)
    assert "weights" in batch and batch["weights"].shape[0] == 8
