# import the necessary packages
from keras.utils import np_utils
import numpy as np
import h5py

class HDF5DatasetGenerator:

    def __init__(self, dbPath, batchSize, preprocessors=None, aug=None, binarize=True, classes=2):
        self.batchSize = batchSize
        self.preprocessors = preprocessors
        self.aug = aug
        self.binarize = binarize
        self.classes = classes
        self.db = h5py.File(dbPath)
        self.numImages = self.db["labels"].shape[0]
    
    def generator(self, passes=np.inf):
        # epochs = 0
        # while epochs < passes:
        for i in np.arange(0, self.numImages, self.batchSize):
            images = self.db["images"][i: i + self.batchSize]
            labels = self.db["labels"][i: i + self.batchSize]
            len_labels = self.db["len_labels"][i: i + self.batchSize]
            if self.binarize:
                labels = np_utils.to_categorical(labels, self.classes)
            if self.preprocessors is not None:
                procImages = []
                for image in images:
                    for p in self.preprocessors:
                        image = p.preprocess(image)
                    procImages.append(image)
                images = np.array(procImages)
            if self.aug is not None:
                (images, labels, len_labels) = next(self.aug.flow(images, labels, len_labels, batch_size=self.batchSize))
            yield (images, labels, len_labels)
            # epochs += 1
    
    def close(self):
        self.db.close()