#!/usr/bin/env python
# -*- coding: utf-8 -*-
from torch.utils.data import Dataset
from pathlib import Path
import numpy as np
from PIL import Image
import torch
import torchvision.transforms.functional as VF
import bob.io.base

def get_file_lists(data_path):
    data_path = Path(data_path)
    image_file_names = np.array(sorted(list(data_path.glob('*'))))
    return image_file_names

class ImageFolderInference(Dataset):
    """
    Generic ImageFolder containing images for inference
    
    Parameters
    ----------
    path : str
        full path to root of dataset
    
    """
    def __init__(self, path, transform = None):
        self.transform = transform
        self.img_file_list = get_file_lists(path)

    def __len__(self):
        """
        Returns
        -------
        int
            size of the dataset
        """
        return len(self.img_file_list)
    
    def __getitem__(self,index):
        """
        Parameters
        ----------
        index : int
        
        Returns
        -------
        list
            dataitem [img_name, img]
        """
        img_path = self.img_file_list[index]
        img_name = img_path.name
        img = Image.open(img_path).convert(mode='RGB')
    
        sample = [img]
        
        if self.transform :
            sample = self.transform(*sample)
        
        sample.insert(0,img_name)
        
        return sample
