"""Synthetic datasets used for zero-cost proxy evaluation."""

from __future__ import annotations

import torch
from torch.utils.data import Dataset


class RandomImageDataset(Dataset):
    """Simple random tensors to satisfy Auto-Prox dataloader contract."""

    def __init__(self, num_samples: int, img_size: int, num_classes: int, in_chans: int = 3):
        self.num_samples = int(num_samples)
        self.img_size = int(img_size)
        self.num_classes = int(num_classes)
        self.in_chans = int(in_chans)

    def __len__(self) -> int:  # pragma: no cover - trivial
        return self.num_samples

    def __getitem__(self, idx):  # pragma: no cover - trivial
        x = torch.rand(self.in_chans, self.img_size, self.img_size)
        y = torch.randint(low=0, high=self.num_classes, size=(1,), dtype=torch.long).item()
        return x, y, torch.zeros(1)

