# coding:utf8

"""
dataloader of DA dataset
"""

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


def generate_label(directory):
	label_name = [i for i in os.listdir(directory) if ".DS_Store" not in i]
	label_dict = {k: v for v, k in enumerate(label_name)}
	return label_dict


class OfficeDataSet(Dataset):
	"""
	Offcei31 dataset
	"""
	def __init__(self, data_path, transformer=None, custom_transformer=None, k=2):
		super(OfficeDataSet, self).__init__()
		self.data_path = data_path
		self.transformer = transformer  # this transformer is generated by torchvision, for reproducing the paper result only
		self.custom_transformer = custom_transformer
		self.k = k  # k represents for num of image transformation
		self.label = generate_label(self.data_path)
		self._parse_directory()
		
	def _load_image(self, img_id):
		img = Image.open(img_id).convert("RGB")
		return img

	def _parse_directory(self):
		self.image_path = [os.path.join(root, file) for root, _, filenames in os.walk(self.data_path) for file in filenames]
		self.label_path = [root.split("/")[-1] for root, _, filenames in os.walk(self.data_path) for file in filenames]
	
	def __getitem__(self, idx):
		img_id = self.image_path[idx]
		label_id = self.label_path[idx]
		img = self._load_image(img_id)
		if self.transformer:
			img = list(map(lambda x: self.transformer(x), [img for i in range(self.k)]))
		img = list(map(lambda x: x.view(1, x.size(0), x.size(1), x.size(2)), img))

		label = [self.label[label_id] for i in range(self.k)]

		return torch.cat(img, 0), torch.tensor(label)

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