import os
import pandas as pd
import h5py

import torch
from torch.utils.data import Dataset

def load_feature(path):
    with h5py.File(path, 'r') as f:
        X = f['X'][()]
    X = torch.FloatTensor(X)
    return X

class VoxelDataset(Dataset):

    def __init__(self, label_file, root_dir):
        """
        Args:
            label_file (string): Path to the csv file with labels.
            root_dir (string): Directory with all the voxel data.
            transform (callable, optional): Optional transform to be applied
                on a sample.
        """
        self.labels = pd.read_csv(label_file)
        self.root_dir = root_dir

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

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

        voxel_name = os.path.join(self.root_dir,
                                self.labels.iloc[idx, 0]) + '.h5'
        voxel = load_feature(voxel_name)
        label = self.labels.iloc[idx, 1]
        sample = {'voxel': voxel, 'label': label}

        return sample