import pytest
import numpy as np
import random
import torch
from ding.data.level_replay.level_sampler import LevelSampler


@pytest.mark.unittest
def test_level_sampler():
    num_seeds = 500
    obs_shape = [3, 64, 64]
    action_shape = 15
    collector_env_num = 16
    level_replay_dict = dict(
        strategy='min_margin',
        score_transform='rank',
        temperature=0.1,
    )
    N = 10
    collector_sample_length = 160

    train_seeds = [i for i in range(num_seeds)]
    level_sampler = LevelSampler(train_seeds, obs_shape, action_shape, collector_env_num, level_replay_dict)

    value = torch.randn(collector_sample_length)
    reward = torch.randn(collector_sample_length)
    adv = torch.randn(collector_sample_length)
    done = torch.randn(collector_sample_length)
    logit = torch.randn(collector_sample_length, N)
    seeds = [random.randint(0, num_seeds) for i in range(collector_env_num)]
    all_seeds = torch.Tensor(
        [seeds[i] for i in range(collector_env_num) for j in range(int(collector_sample_length / collector_env_num))]
    )

    train_data = {'value': value, 'reward': reward, 'adv': adv, 'done': done, 'logit': logit, 'seed': all_seeds}
    level_sampler.update_with_rollouts(train_data, collector_env_num)
    sample_seed = level_sampler.sample()
    assert isinstance(sample_seed, int)