import math
import numpy as np

from delira.data_loading import AbstractDataset


class DummyDataset(AbstractDataset):
    def __init__(self, length=600, class_weights=None):
        super().__init__(None, None)

        if class_weights is None:
            class_weights = [0.5, 0.3, 0.2]
        assert math.isclose(sum(class_weights), 1)

        self._data = [np.random.rand(1, 28, 28) for i in range(length)]
        _labels = []
        for idx, weight in enumerate(class_weights):
            _labels += [idx] * int(length * weight)

        self._labels = _labels

    def __getitem__(self, index):
        return {"data": self._data[index], "label": self._labels[index]}

    def __len__(self):
        return len(self._data)
