from torch.utils.data import Dataset

from .data_utils import load_image

from .parsers import BaseParser
from .data_utils import LabelTransform
from .perceptual_transforms import PerceptualFeatureMapTransform


class EPUDataset(Dataset):
    def __init__(self,
                 data_parser: BaseParser,
                 perceptual_transform: PerceptualFeatureMapTransform = None,
                 label_transform: LabelTransform = None,
                 **kwargs):

        self.data_parser = data_parser
        self.transform = perceptual_transform
        self.label_transform = label_transform
        self.kwargs = kwargs

        self.image_paths = self.data_parser.image_filenames
        self.labels = self.data_parser.labels

    def __len__(self):
        assert len(self.image_paths) == len(self.labels), "Mismatch in image paths and labels"
        return len(self.image_paths)

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        label = self.labels[idx]

        img = load_image(img_path)
        if self.transform is not None:
            img = self.transform(img)
        if self.label_transform is not None:
            label = self.label_transform(label)
        return img, label