class Cifar10(Dataset):
    mode: Any = ...
    backend: Any = ...
    data_file: Any = ...
    transform: Any = ...
    dtype: Any = ...

from __future__ import annotations

import numpy as np
import numpy.typing as npt

from typing import overload, TypeVar, Any, Sequence, Callable, Tuple, Union


class Tensor: ...


class PILImage: ...


_DataT = TypeVar("_DataT", Tensor, PILImage, npt.NDArray[Any])


# ---------- transform ----------
_TransformTensor = Callable[[Tensor], Tensor]
_TransformTensors = Callable[[Tuple[Tensor, ...]], Tuple[Tensor, ...]]
_TransformPILImage = Callable[[PILImage], PILImage]
_TransformPILImages = Callable[[Tuple[PILImage, ...]], Tuple[PILImage, ...]]
_TransformNDArray = Callable[[npt.NDArray[Any]], npt.NDArray[Any]]
_TransformNDArrays = Callable[
    [Tuple[npt.NDArray[Any], ...]], Tuple[npt.NDArray[Any], ...]
]

_Transform = Union[
    _TransformTensor,
    _TransformTensors,
    _TransformPILImage,
    _TransformPILImages,
    _TransformNDArray,
    _TransformNDArrays,
]
# ---------- transform end ----------


class BaseTransform:
    def __init__(self, keys=None):
        self.keys = keys

    @overload
    def __call__(self, inputs: _DataT) -> _DataT: ...

    @overload
    def __call__(self, inputs: tuple[_DataT, ...]) -> tuple[_DataT, ...]: ...

    def __call__(self, inputs):
        if isinstance(inputs, (Tensor, PILImage, np.ndarray)):
            return inputs
        return tuple(inputs)


class Resize(BaseTransform):
    def __init__(self, size, interpolation="bilinear", keys=None): ...

    @overload
    def __call__(self, inputs: _DataT) -> _DataT: ...

    @overload
    def __call__(self, inputs: tuple[_DataT, ...]) -> tuple[_DataT, ...]: ...

    def __call__(self, inputs):
        if isinstance(inputs, (Tensor, PILImage, np.ndarray)):
            return inputs
        return tuple(inputs)


class Compose:
    transforms: Sequence[BaseTransform]

    def __init__(self, transforms: Sequence[BaseTransform]) -> None:
        self.transforms = transforms

    @overload
    def __call__(self, data: _DataT) -> _DataT: ...

    @overload
    def __call__(self, data: tuple[_DataT, ...]) -> tuple[_DataT, ...]: ...

    def __call__(self, data):
        if isinstance(data, (Tensor, PILImage, np.ndarray)):
            return data
        return tuple(data)


class Cifar10:
    transform: _Transform

    def __init__(self, transform: _Transform) -> None:
        self.transform = transform


t1 = Resize(1)
cifar = Cifar10(transform=t1) # ok

t2 = Compose([Resize(1), Resize(2)])
cifar = Cifar10(transform=t2) # ok


def trans_pass(a: Tensor) -> Tensor:
    return Tensor()


def trans_fail(a: Tensor) -> tuple[Tensor, Tensor]:
    return Tensor(), Tensor()


cifar = Cifar10(transform=trans_pass) # ok
cifar = Cifar10(transform=trans_fail) # This should fail