import os
import numpy as np
from PIL import Image
import torch
from torch.utils.data import DataLoader, Subset
from torchvision import transforms

class TinyImages(torch.utils.data.Dataset):

	def __init__(self, root, transform):

		super(TinyImages, self).__init__()

		self.data = np.load(os.path.join(root, 'tinyimages80m', '300K_random_images.npy'))
		self.transform = transform

		print("TinyImages Contain {} images".format(len(self.data)))

	def __getitem__(self, index):
		img = self.data[index]
		img = Image.fromarray(np.uint8(img))
		if self.transform is not None:
			img = self.transform(img)

		return img, -1  # -1 is the class

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


def tinyimages300k_dataloaders(num_samples=300000, train_batch_size=64, num_workers=8, data_root_path='/ssd1/haotao/datasets'):

	num_samples = int(num_samples)

	data_dir = os.path.join(data_root_path)

	train_transform = transforms.Compose([
		transforms.RandomCrop(32, padding=4),
		transforms.RandomHorizontalFlip(),
		transforms.ToTensor(),
	])

	train_set = Subset(TinyImages(data_dir, train=True, transform=train_transform, download=True), list(range(num_samples)))

	train_loader = DataLoader(train_set, batch_size=train_batch_size, shuffle=True, num_workers=num_workers,
								drop_last=True, pin_memory=True)

	return train_loader
	