import torch
import torchvision.transforms as transforms
from torch.utils.data import Dataset
import glob
from PIL import Image
import random


class SUN397EncodableDataset(Dataset):
    """SUN397 encodable dataset class"""

    def __init__(self, train=True):
        super().__init__()
        path = 'data/SUN397/train/*/*.jpg' if train else 'data/SUN397/test/*/*.jpg'
        self.data = list(glob.glob(path))
        random.shuffle(self.data)
        cats = list(set([path.split("/")[3] for path in self.data]))
        cats.sort()
        self.labels = torch.LongTensor([cats.index(path.split("/")[3]) for path in self.data])
        self.preprocessor = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        if len(self.encoded_data) == 0:
            return self.preprocessor(Image.open(self.data[idx]).convert('RGB')), self.labels[idx]
        return self.encoded_data[idx], self.labels[idx]

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

    def num_classes(self):
        return int(max(self.labels) + 1)

[(16, 'random.shuffle', 'random.shuffle', False, 'import random\n'), (28, 'torch.is_tensor', 'torch.is_tensor', False, 'import torch\n'), (15, 'glob.glob', 'glob.glob', False, 'import glob\n'), (21, 'torchvision.transforms.Resize', 'transforms.Resize', True, 'import torchvision.transforms as transforms\n'), (22, 'torchvision.transforms.ToTensor', 'transforms.ToTensor', True, 'import torchvision.transforms as transforms\n'), (23, 'torchvision.transforms.Normalize', 'transforms.Normalize', True, 'import torchvision.transforms as transforms\n'), (25, 'torch.cuda.is_available', 'torch.cuda.is_available', False, 'import torch\n'), (32, 'PIL.Image.open', 'Image.open', False, 'from PIL import Image\n')]