import numpy as np
import torch
from logger import Logger
LOGGER = Logger().logger()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
import glob
import os
from PIL import Image
import random

#Author: @meetdoshi
class LoadData(Dataset):
    def __init__(self, data, pair=False, type_data="train"):
        self.data = data
        self.pair = pair
        self.type = type_data
        self.transform = transforms.Compose([
                transforms.Resize(int(256 * 1.12), Image.BICUBIC),
                transforms.RandomCrop(256),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
        ])
        #LOGGER.info(data)
        #LOGGER.info(os.path.join(self.data,f"{type_data}/A"+ "*.jpg"))
        self.X = sorted(glob.glob(os.path.join(self.data, f"{type_data}/A/")+ "*.jpg"))
        self.Y = sorted(glob.glob(os.path.join(self.data, f"{type_data}/B/")+ "*.jpg"))


    def __len__(self):
        return max(len(self.X), len(self.Y))

    def __getitem__(self, idx):
        if self.pair:
            return {'X':self.transform(Image.open(self.X[idx%len(self.X)])), 'Y':self.transform(Image.open(self.Y[idx%len(self.Y)]))}
        else:
            return {'X':self.transform(Image.open(self.X[idx%len(self.X)])), 'Y':self.transform(Image.open(self.Y[random.randint(0, len(self.Y)-1)%len(self.Y)]))}
