# Wraping for the CIFAR dataset

from typing import Callable, Optional

from torchvision.datasets import SUN397 as SUN397_tv
from torchvision.transforms import transforms


class SUN397(SUN397_tv):

    def __init__(
        self,
        root: str,
        train: bool = True,
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        download: bool = False,
    ) -> None:
        assert train is False
        super().__init__(root,
                         transform=transforms.ToTensor()
                         if transform is None else transform,
                         target_transform=target_transform,
                         download=download)

        self.classes_names = self.classes
        self.targets = self._labels

        self.mean, self.std = (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
        if isinstance(self.transform, transforms.Compose):
            self.transform.transforms.append(
                transforms.Normalize(self.mean, self.std))


if __name__ == "__main__":
    dataset = SUN397(root="data", train=False, download=True)
    print(len(dataset), len(dataset.classes_names))
