import d4rl
import gym
import numpy as np

from jax_rl.datasets.dataset import Dataset
from jax_rl.datasets.replay_buffer import Batch


class D4RLDataset(Dataset):
    def __init__(self,
                 env: gym.Env,
                 clip_to_eps: bool = True,
                 eps: float = 1e-5):
        self.dataset = d4rl.qlearning_dataset(env)

        if clip_to_eps:
            lim = 1 - eps
            self.dataset['actions'] = np.clip(self.dataset['actions'], -lim,
                                              lim)

    def sample(self, batch_size: int) -> Batch:
        indx = np.random.randint(len(self.dataset['observations']),
                                 size=batch_size)
        return Batch(observations=self.dataset['observations'][indx],
                     actions=self.dataset['actions'][indx],
                     rewards=self.dataset['rewards'][indx],
                     masks=1.0 -
                     self.dataset['terminals'][indx].astype(np.float32),
                     next_observations=self.dataset['next_observations'][indx])
