import torch
from torchrl.data.replay_buffers import ReplayBuffer, LazyTensorStorage
from torchrl.data.replay_buffers.samplers import SliceSampler, PrioritizedSliceSampler
from tensordict import TensorDict


def test_sampler():
    torch.manual_seed(0)

    sampler = PrioritizedSliceSampler(
        max_capacity=10,
        num_slices=2,
        traj_key="trajectory",
        # end_key="done",
        strict_length=True,
        alpha=1.0,
        beta=1.0,
    )
    trajectory = torch.tensor([3, 0, 1, 1, 1, 2, 2, 2, 3, 3])
    done = torch.tensor([True, True, False, False, True, False, False, True, False, True])
    td = TensorDict({"trajectory": trajectory, "steps": torch.arange(10), "done": done}, [10])

    rb = ReplayBuffer(
        sampler=sampler,
        storage=LazyTensorStorage(10, device=torch.device("cpu")),
        batch_size=6,
    )

    rb.extend(td)
    for i in range(10):
        # preceding_stop_idx in sample(): [1 2 3 5 6 8 9]
        s, info = rb.sample(return_info=True)
        traj = s["trajectory"]
        print("[loop {}] sampled trajectory: {}".format(i, traj))
        print("[loop {}] index {}".format(i, info["index"]))
        assert len(traj.unique())<=2


test_sampler()