import numpy as np
import cv2
from torch.utils.data import Dataset
import torch
import random
IGNORED = ['.DS_Store']


class MRIDataset(Dataset):
    def __init__(self, imgpath, labelpath, preprocessors=None, verbose=-1):
        super(MRIDataset, self).__init__()
        # store the image preprocessor
        self.preprocessors = preprocessors
        self.imgpath = imgpath
        self.labelpath = labelpath

        # if the preprocessors are None, initialize them as an
        # empty list
        if self.preprocessors is None:
            self.preprocessors = []

        self.images = []
        self.masks = []

        for (i, path) in enumerate(self.imgpath):
            image = cv2.imread(path,0)

            if self.preprocessors is not None:
                for p in self.preprocessors:
                    image = p.preprocess(image)

            image = torch.from_numpy(image)
            image = image.unsqueeze(0)

            self.images.append(image)

            if verbose > 0 and i > 0 and (i + 1) % verbose == 0:
                print("[INFO] processed {}/{}".format(i + 1, len(path)))

        for (i, path) in enumerate(self.labelpath):
            label = cv2.imread(path)

            if self.preprocessors is not None:
                for p in self.preprocessors:
                    label = p.preprocess(label)
            label = np.sum(label, axis=2)
            label = label > 0.5
            label = torch.from_numpy(label)
            label = label.unsqueeze(0)

            self.masks.append(label)

            if verbose > 0 and i > 0 and (i + 1) % verbose == 0:
                print("[INFO] processed {}/{}".format(i + 1, len(path)))

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

    def __getitem__(self, idx):

        image = self.images[idx]
        mask = self.masks[idx]

        # Flip image for data augmentation
        if random.random() > 0.5:
            image = torch.flip(image, [0])
            mask = torch.flip(mask, [0])

        return image.float(), mask.float()