import glob
import random
import os

from torch.utils.data import Dataset
from PIL import Image
import numpy as np
import torchvision.transforms as transforms

class ImageDataset(Dataset):
    def __init__(self, dataset_dir, transforms_=None):
        '''
        Construct a dataset with all images from a dir.

        dataset_dir: str. img folder path
        '''
        self.transform = transforms.Compose(transforms_)

        self.files = sorted(glob.glob(os.path.join(dataset_dir) + '/*.jpg'))
        print('files:', len(self.files))

    def __getitem__(self, index):
        img = Image.open(self.files[index % len(self.files)])
        img = img.convert("RGB")
        item = self.transform(img)

        return item

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


# class PairedImageDataset(Dataset):
#     def __init__(self, dataset_dir, soft_data_dir, mode='A2B', transforms_=None):
#         '''
#         Construct a dataset with all images from a dir.

#         dataset: str. dataset name
#         style: str. 'A2B' or 'B2A'
#         '''
#         self.transform = transforms.Compose(transforms_)
        
#         if mode in ['A2B']:
#             path_A = os.path.join(dataset_dir, 'train', 'A')
#             path_B = os.path.join(soft_data_dir, 'B')
#         else:
#             path_A = os.path.join(soft_data_dir, 'A')
#             path_B = os.path.join(dataset_dir, 'train', 'B')
#         self.files_A = sorted(glob.glob(path_A + '/*.*'))
#         self.files_B = sorted(glob.glob(path_B + '/*.*'))
#         print('files_A:', len(self.files_A))
#         print('files_B:', len(self.files_B))
#         assert len(self.files_A) == len(self.files_B)

#     def __getitem__(self, index):
#         img_A = Image.open(self.files_A[index % len(self.files_A)])
#         img_A = img_A.convert("RGB")
#         img_A= np.asarray(img_A) # PIL.Image to np.ndarray
#         img_A = np.flip(img_A, axis=1) # data augumentation: horrizental flip
#         img_A = Image.fromarray(np.uint8(img_A)) # np.ndarray to PIL.Image
#         item_A = self.transform(img_A)

#         img_B = Image.open(self.files_B[index % len(self.files_B)])
#         img_B = img_B.convert("RGB")
#         img_B= np.asarray(img_B) # PIL.Image to np.ndarray
#         img_B = np.flip(img_B, axis=1) # data augumentation: horrizental flip
#         img_B = Image.fromarray(np.uint8(img_B)) # np.ndarray to PIL.Image
#         item_B = self.transform(img_B)
#         return {'A': item_A, 'B': item_B}

#     def __len__(self):
#         return len(self.files_A)
