"""Deterministic fake dataset for the tiny AtmoRep training pipeline."""

import torch
from torch.utils.data import Dataset

from .tiny_atmorep import TinyAtmoRepConfig


class FakeAtmoRepDataset(Dataset):
    def __init__(self, config: TinyAtmoRepConfig, samples: int, seed: int, mask_fraction: float) -> None:
        self.config, self.samples, self.seed = config, samples, seed
        self.num_tokens = (config.input_shape[0] // config.patch_shape[0]) * (config.input_shape[2] // config.patch_shape[1]) * (config.input_shape[3] // config.patch_shape[2])
        self.mask_count = max(1, round(self.num_tokens * mask_fraction))

    def __len__(self):
        return self.samples

    def __getitem__(self, index):
        generator = torch.Generator().manual_seed(self.seed + index)
        fields = torch.randn(self.config.input_shape, generator=generator)
        mask = torch.zeros(self.num_tokens, dtype=torch.bool)
        mask[torch.randperm(self.num_tokens, generator=generator)[:self.mask_count]] = True
        return fields, mask