from os import listdir

from os.path import join

import PIL.Image

import matplotlib.pyplot as plt

import numpy as np

import torch.utils.data

from PIL.Image import open

from torch.nn.functional import one_hot

from torch.utils.data import Dataset

from torchvision import transforms

from perturbation.model import ModelOfNaturalVariation

import data.augmentation as aug

from utils.mask_generator import generate_a_mask

class KvasirClassificationDataset(Dataset):

    """

    Dataset class that fetches images with the associated pathological class labels for use in Pretraining

    """

    def __init__(self, path):

        super(KvasirClassificationDataset, self).__init__()

        self.path = join(path, "labeled-images/lower-gi-tract/pathological-findings")

        self.label_names = listdir(self.path)

        self.num_classes = len(self.label_names)

        self.fname_class_dict = {}

        i = 0

        self.class_weights = np.zeros(self.num_classes)

        for i, label in enumerate(self.label_names):

            class_path = join(self.path, label)

            for fname in listdir(class_path):

                self.class_weights[i] += 1

                self.fname_class_dict[fname] = label

        self.index_dict = dict(zip(self.label_names, range(self.num_classes)))

        self.common_transforms = transforms.Compose([transforms.Resize((400, 400)),

                                                     transforms.ToTensor()

                                                     ])

    def __len__(self):

        # return 256  # for debugging

        return len(self.fname_class_dict)

    def __getitem__(self, item):

        fname, label = list(self.fname_class_dict.items())[item]

        onehot = one_hot(torch.tensor(self.index_dict[label]), num_classes=self.num_classes)

        image = open(join(join(self.path, label), fname)).convert("RGB")

        # print(list(image.getdata()))

        # input()

        image = self.common_transforms(open(join(join(self.path, label), fname)).convert("RGB"))

        return image, onehot.float(), fname

class KvasirSegmentationDataset(Dataset):

    """

        Dataset class that fetches images with the associated segmentation mask.

        Employs "vanilla" augmentations

    """

    def __init__(self, path, split="train", augment=False):

        super(KvasirSegmentationDataset, self).__init__()

        self.path = join(path, "segmented-images/")

        self.fnames = listdir(join(self.path, "images"))

        self.common_transforms = aug.pipeline_tranforms()

        self.pixeltrans = aug.albumentation_pixelwise_transforms()

        self.segtrans = aug.albumentation_pixelwise_transforms()

        # deterministic partition

        self.split = split

        train_size = int(len(self.fnames) * 0.8)

        val_size = (len(self.fnames) - train_size) // 2

        test_size = len(self.fnames) - train_size - val_size

        self.augment = augment

        self.fnames_train = self.fnames[:train_size]

        self.fnames_val = self.fnames[train_size:train_size + val_size]

        self.fnames_test = self.fnames[train_size + val_size:]

        self.split_fnames = None  # iterable for selected split

        if self.split == "train":

            self.size = train_size

            self.split_fnames = self.fnames_train

        elif self.split == "val":

            self.size = val_size

            self.split_fnames = self.fnames_val

        elif self.split == "test":

            self.size = test_size

            self.split_fnames = self.fnames_test

        else:

            raise ValueError("Choices are train/val/test")

    def __len__(self):

        return self.size

    def __getitem__(self, index):

        image = np.array(open(join(self.path, "images/", self.split_fnames[index])).convert("RGB"))

        mask = np.array(open(join(self.path, "masks/", self.split_fnames[index])).convert("L"))

        if self.split == "train" and self.augment == True:

            transformed = self.pixeltrans(image=image)

            image = transformed["image"]

            segtransformed = self.segtrans(image=image, mask=mask)

            image, mask = segtransformed["image"], segtransformed["mask"]

        image = self.common_transforms(PIL.Image.fromarray(image))

        mask = self.common_transforms(PIL.Image.fromarray(mask))

        mask = (mask > 0.5).float()

        return image, mask, self.split_fnames[index]

class KvasirMNVset(KvasirSegmentationDataset):

    def __init__(self, path, split, inpaint=False):

        super(KvasirMNVset, self).__init__(path, split, augment=False)

        self.mnv = ModelOfNaturalVariation(1, use_inpainter=True)

        self.p = 0.5

    def __getitem__(self, index):

        image = np.array(open(join(self.path, "images/", self.split_fnames[index])).convert("RGB"))

        mask = np.array(open(join(self.path, "masks/", self.split_fnames[index])).convert("L"))

        image = self.common_transforms(PIL.Image.fromarray(image))

        mask = self.common_transforms(PIL.Image.fromarray(mask))

        mask = (mask > 0.5).float()

        flag = False

        if self.split == "train" and np.random.rand() < self.p:

            flag = True

            image, mask = self.mnv(image.unsqueeze(0), mask.unsqueeze(0))

            image = image.squeeze()

            mask = mask.squeeze(0)  # todo make this less ugly

            # plt.imshow(image.T.cpu().numpy())

            # plt.show()

        return image, mask, self.split_fnames[index], flag

    def set_prob(self, prob):

        self.p = prob

class KvasirInpaintingDataset(Dataset):

    def __init__(self, path, split="train"):

        super(KvasirInpaintingDataset, self).__init__()

        self.path = join(path, "segmented-images/")

        self.fnames = listdir(join(self.path, "images"))

        self.common_transforms = transforms.Compose([transforms.Resize((400, 400)),

                                                     transforms.ToTensor()

                                                     ])

        self.split = split

        train_size = int(len(self.fnames) * 0.8)

        val_size = (len(self.fnames) - train_size) // 2

        test_size = len(self.fnames) - train_size - val_size

        self.fnames_train = self.fnames[:train_size]

        self.fnames_val = self.fnames[train_size:train_size + val_size]

        self.fnames_test = self.fnames[train_size + val_size:]

        self.split_fnames = None  # iterable for selected split

        if self.split == "train":

            self.size = train_size

            self.split_fnames = self.fnames_train

        elif self.split == "val":

            self.size = val_size

            self.split_fnames = self.fnames_val

        elif self.split == "test":

            self.size = test_size

            self.split_fnames = self.fnames_test

        else:

            raise ValueError("Choices are train/val/test")

    def __len__(self):

        return len(self.split_fnames)

    def __getitem__(self, index):

        image = self.common_transforms(

            open(join(join(self.path, "images/"), self.split_fnames[index])).convert("RGB"))

        mask = self.common_transforms(

            open(join(join(self.path, "masks/"), self.split_fnames[index])).convert("L"))

        mask = (mask > 0.5).float()

        part = mask * image

        masked_image = image - part

        return image, mask, masked_image, part, self.split_fnames[index]

class KvasirSyntheticDataset(Dataset):

    def __init__(self, path, split="train"):

        super(KvasirSyntheticDataset, self).__init__()

        self.path = join(path, "unlabeled-images")

        self.fnames = listdir(join(self.path, "images"))

        self.common_transforms = aug.pipeline_tranforms()

        self.split = split

        train_size = int(len(self.fnames) * 0.8)

        val_size = (len(self.fnames) - train_size) // 2

        test_size = len(self.fnames) - train_size - val_size

        self.fnames_train = self.fnames[:train_size]

        self.fnames_val = self.fnames[train_size:train_size + val_size]

        self.fnames_test = self.fnames[train_size + val_size:]

        self.split_fnames = None  # iterable for selected split

        if self.split == "train":

            self.size = train_size

            self.split_fnames = self.fnames_train

        elif self.split == "val":

            self.size = val_size

            self.split_fnames = self.fnames_val

        elif self.split == "test":

            self.size = test_size

            self.split_fnames = self.fnames_test

        else:

            raise ValueError("Choices are train/val/test")

        print("loading mnv")

        self.mnv = ModelOfNaturalVariation(0, use_inpainter=True).to("cuda")

        print("mnv loaded")

    def __len__(self):

        return len(self.split_fnames)

        # return 10  # debug

    def __getitem__(self, index):

        print(f"getting {index}")

        image = np.array(open(join(self.path, "images/", self.split_fnames[index])).convert("RGB"))

        mask = np.zeros_like(image)

        image = self.common_transforms(PIL.Image.fromarray(image))

        mask = self.common_transforms(PIL.Image.fromarray(mask))

        image, mask = self.mnv(image.unsqueeze(0), mask.unsqueeze(0))

        image = image.squeeze()

        mask = mask.squeeze(0)

        return image, mask, self.split_fnames[index]

def test_KvasirSegmentationDataset():

    dataset = KvasirSegmentationDataset("Datasets/HyperKvasir", split="test")

    for x, y, fname in torch.utils.data.DataLoader(dataset):

        plt.imshow(x.squeeze().T)

        # plt.imshow(y.squeeze().T, alpha=0.5)

        plt.show()

        assert isinstance(x, torch.Tensor)

        assert isinstance(y, torch.Tensor)

    print("Segmentation evaluation passed")

def test_KvasirClassificationDataset():

    dataset = KvasirClassificationDataset("Datasets/HyperKvasir")

    for x, y, fname in torch.utils.data.DataLoader(dataset):

        assert isinstance(x, torch.Tensor)

        assert isinstance(y, torch.Tensor)

    print("Classification evaluation passed")

if __name__ == '__main__':

    test_KvasirSegmentationDataset()

    # test_KvasirClassificationDataset()