"""Indexed deterministic cubed-sphere samples for training and validation."""

from torch.utils.data import Dataset
from .fake_data import make_fake_batch


class FakeCubeSphereDataset(Dataset):
    def __init__(self, samples=64, seed=7, channels=2, height=8, width=8):
        self.samples, self.seed = samples, seed
        self.shape = (channels, 6, height, width)

    def __len__(self):
        return self.samples

    def __getitem__(self, index):
        return make_fake_batch(batch=1, channels=self.shape[0], faces=6, height=self.shape[2], width=self.shape[3], seed=self.seed + index)[0]