import os
import torch
from torch.utils.data import IterableDataset, Dataset
from torchvision.transforms import v2
from PIL import Image
import random
from abc import ABC

class ImageDataset(ABC):
    """
    Class, to be used as base dataset class, which models a dataset of images, where images
    are loaded from file. All images are transformed to grayscale images and
    represented as torch tensors of the specified data type (by default float64)
    """
    def __init__(self, root_path: str, dtype: torch.dtype = torch.float64) -> None:
        """
        Initialisation of class ImageDataset

        :param root_path: Path to the directory containing the images
        :param dtype: Datatype which is used to represent images in terms of tensors.
        """
        self._root_path = root_path
        self._image_list = os.listdir(root_path)
        self._transform = v2.Compose([v2.ToImage(), v2.ToDtype(dtype, scale=True),
                                      v2.Grayscale(num_output_channels=1)])

    def _load_image(self, image_name: str) -> Image.Image:
        return Image.open(os.path.join(self._root_path, image_name))

class TestImageDataset(ImageDataset, Dataset):
    """
    Class, inheriting from ImageDataset and torch.utils.data.Dataset, which is used to model
    a dataset specifically for testing.
    """

    def __init__(self, root_path: str, dtype: torch.dtype = torch.float64):
        ImageDataset.__init__(self, root_path, dtype)
        Dataset.__init__(self)

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

    def __getitem__(self, idx: int):
        image_name = self._image_list[idx]
        image = self._load_image(image_name)
        return self._transform(image).unsqueeze(dim=0)

class TrainingImageDataset(ImageDataset, IterableDataset):
    """
    Class, inheriting from ImageDataset and torch.utils.data.IterableDataset, modelling an image dataset
    for the purpose of training: Images are loaded randomly from file - sampling with replacement. Note that
    crops will not be generated by methods of this class - use a data loader and a collate function instead.

    """
    def __init__(self, root_path: str, dtype: torch.dtype = torch.float64) -> None:
        """
        Initialisation of an object of class TrainingImageDataset

        :param root_path: Path to image directory
        :param dtype: torch datatype the images are represented with
        """
        ImageDataset.__init__(self, root_path, dtype)
        Dataset.__init__(self)

    def __iter__(self) -> torch.Tensor:
        while True:
            image_name = random.choice(self._image_list)
            image = self._load_image(image_name)
            yield self._transform(image).unsqueeze(dim=0)
