# SPDX-License-Identifier: MIT

try:
    from extra_data import H5File
    from extra_data.exceptions import FileStructureError
    extra_data_installed = True
except ImportError:
    extra_data_installed = False

__all__ = ['KaraboFile', 'extra_data_installed']


class KaraboFile:
    def __init__(self, filename, source_ind=0):
        if not extra_data_installed:
            raise IOError('extra_data is required to load karabo h5 files')
        try:
            self.f = H5File(filename)
        except FileStructureError:
            raise IOError('This hdf5 file is not a generated by karabo')
        self.series_max = len(self.f.train_ids)
        self.sources = [s for s in self.f.instrument_sources if "daqOutput" in s]
        self.current_source = self.sources[source_ind]

    def get_image(self, ind):
        sel = self.f.select(self.current_source, 'data.image.pixels')
        tid, data = sel.train_from_index(ind)
        return data[self.current_source]['data.image.pixels'][::-1]
