import torch
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt

import os
import pandas as pd
from torchvision.io import read_image


def show_samples():
    training_data = datasets.FashionMNIST(
        root="data",
        train=True,
        download=False,
        transform=ToTensor()
    )
    test_data = datasets.FashionMNIST(
        root="data",
        train=False,
        download=False,
        transform=ToTensor()
    )

    labels_map = {
        0: "T-Shirt",
        1: "Trouser",
        2: "Pullover",
        3: "Dress",
        4: "Coat",
        5: "Sandal",
        6: "Shirt",
        7: "Sneaker",
        8: "Bag",
        9: "Ankle Boot",
    }
    figure = plt.figure(figsize=(8, 8))
    cols, rows = 3, 3
    for i in range(1, cols * rows + 1):
        sample_idx = torch.randint(len(training_data), size=(1,)).item()
        img, label = training_data[sample_idx]
        figure.add_subplot(rows, cols, i)
        plt.title(labels_map[label])
        plt.axis("off")
        plt.imshow(img.squeeze(), cmap="gray")
    plt.show()
    pass


class CustomImageDataset(Dataset):
    """
    It should be clear that the dataset is stored in the disk rather than in the memory,
    because the dataset is usually huge.

    This class can only load image datasets, the core method called in this class is
    `read_image`, which can load an image into memory according to a given image full file
    name.

    A dataset class must satisfy the following requirements:
    - inherit from torch.utils.data.Dataset
    - realize/overload __init__(), __len__(), and __getitem()__ method. The return data
    type of __getitem__() must be torch.Tensor

    Ultrasound RF signals are usually stored in binary format in the disk, and the value
    can be positive or negative, so when loaded to memory, it can be read as a
    numpy.ndarray first, and then convert to a Tensor. It should also be normalized.
    """

    def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
        """
        construction method.

        :param annotations_file: a csv file which contains all samples file name and
            corresponding labels value.
        :param img_dir: the path where samples files are stored.
        :param transform:
        :param target_transform:
        """
        self.img_labels = pd.read_csv(annotations_file)
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform

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

    def __getitem__(self, idx):
        """
        return the corresponding sample and label according to the given index.

        :param idx:
        :return:
        """
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
        image = read_image(img_path)
        label = self.img_labels.iloc[idx, 1]
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            label = self.target_transform(image)
        return image, label


def main():
    """
    create a custom dataset

    :return:
    """

    pass


if __name__ == '__main__':
    main()