import os
from pathlib import Path
from timm.data.readers.reader import Reader


class AnnotationReader(Reader):

    def __init__(
            self,
            root,
            images, 
            labels):
        super().__init__()

        self.root = root
        self.samples = list(zip(images, labels))
        
        if len(self.samples) == 0:
            raise RuntimeError(
                f'Found 0 images')

    def __getitem__(self, index):
        _path, target = self.samples[index]
        path = Path(self.root, _path)
        return open(path, 'rb'), target

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

    def _filename(self, index, basename=False, absolute=False):
        filename = self.samples[index][0]
        if basename:
            filename = os.path.basename(filename)
        elif not absolute:
            filename = os.path.relpath(filename, self.root)
        return filename