import os
import torch
import numpy as np
from torch.utils.data.dataset import Dataset
from collections import defaultdict

import nibabel as nib



class UnetDataset(Dataset):
	def __init__(self, location_metadata, list_subjects ,dirpath_gm, num_voxels):
		super(UnetDataset, self).__init__()
		self.location_metadata = location_metadata
		self.list_subjects = list_subjects
		self.dirpath_gm = dirpath_gm
		self.num_voxels = num_voxels

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

	def __getitem__(self, index):
		current_subject = self.list_subjects[index]

		current_row_of_interrst = self.location_metadata.loc[current_subject]

		subject_age = current_row_of_interrst.Age / 100
		subject_gender = current_row_of_interrst.Gender
		subject_name = current_subject

		temporar_object = nib.load(self.dirpath_gm+'/'+subject_name)
		temporar_data = temporar_object.get_data()
		temporar_object.uncache()

		#x_training_data = np.expand_dims(temporar_data,axis=0)
		x_training_data = temporar_data

		return 	x_training_data,subject_age,subject_gender