from typing import Iterable, override

import numpy as np

from ai_parade._toolkit.data.ImageInput import ImageInput

from .Dataset import Dataset, DatasetInfo, Output_unbatched


class Random(Dataset[None]):
	"""Dataset with random data of the ImageInput specifications"""

	def __init__(
		self,
		image_input: ImageInput | None = None,
		*,
		seed: int | None = None,
		size: int = 1000,
	):
		super().__init__()
		self._image_input = image_input
		if image_input is not None:
			self._batch_size = image_input.batchSize
		self._seed = seed
		self._size = size

	@override
	def walk(self) -> Iterable[None]:
		for _ in range(self._size):
			yield None

	@override
	def parse_data(
		self, iterable: Iterable[None], should_load: bool
	) -> Iterable[Output_unbatched]:
		assert self._image_input is not None, (
			"The Random dataset was constructed without an ImageInput or preprocess call"
		)
		rng = np.random.RandomState(self._seed)
		shape = self._image_input.shape[1:]
		for _ in iterable:
			match self._image_input.dataType:
				case ImageInput.DataType.float32:
					data = rng.rand(*shape).astype(np.float32)
				case ImageInput.DataType.uint8:
					data = rng.randint(
						0,
						255,
						size=shape,
						dtype=np.uint8,
					)
			yield (
				data,
				{},
				DatasetInfo(
					source_data_shape=shape,
					data_source="Random",
					data_shape=shape,
				),
			)

	@override
	def preprocess(self, image_input: ImageInput):
		self._image_input = image_input
		self._batch_size = image_input.batchSize
		return self
