from dataclasses import dataclass
from enum import StrEnum
from typing import List, Optional


@dataclass(frozen=True)
class ImageInput:
	"""
	Represents the input specification for a CV model.
	"""

	batchSize: int
	"""Number of images in a single batch."""

	height: int
	"""Height of the input image."""

	width: int
	"""Width of the input image."""

	channelOrder: "ImageInput.ChannelOrder"
	"""Order of color channels in the image."""

	dataOrder: "ImageInput.DataOrder"
	"""Order of data dimensions."""

	dataType: "ImageInput.DataType"
	"""Data type of the image pixels."""

	means: Optional[tuple[float, float, float] | tuple[float, float, float, float]] = (
		None
	)
	"""
	Mean values for normalization, per channel.
	In [0, 255] range for integer data (`dataType=uint8`).
	In [0, 1] range for float data, (`dataType=float32` or `dataType=float16`).
	"""

	stds: Optional[tuple[float, float, float] | tuple[float, float, float, float]] = (
		None
	)
	"""
	Standard deviation values for normalization, per channel.
	In [0, 255] range for integer data (`dataType=uint8`).
	In [0, 1] range for float data, (`dataType=float32` or `dataType=float16`).
	"""

	@property
	def channels(self) -> int:
		"""
		Number of channels in the image.
		"""
		return (
			4
			if self.channelOrder
			in {ImageInput.ChannelOrder.RGBA, ImageInput.ChannelOrder.BGRA}
			else 3
		)

	@property
	def shape(self) -> List[int]:
		"""
		Shape of the input image, including batch size
		"""
		return (
			[self.batchSize, self.channels, self.height, self.width]
			if self.dataOrder == ImageInput.DataOrder.NCHW
			else [self.batchSize, self.height, self.width, self.channels]
		)

	class DataOrder(StrEnum):
		"""Enumeration for data order formats."""

		NHWC = "NHWC"
		NCHW = "NCHW"

	class DataType(StrEnum):
		"""Enumeration for data types."""

		uint8 = "uint8"
		float32 = "float32"
		float16 = "float16"

	class ChannelOrder(StrEnum):
		"""Enumeration for channel order types."""

		RGB = "RGB"
		BGR = "BGR"
		RGBA = "RGBA"
		BGRA = "BGRA"
