import os
import numpy as np
import cv2
import albumentations
from PIL import Image
from torch.utils.data import Dataset

from taming.data.sflckr import SegmentationBase # for examples included in repo


class Examples(SegmentationBase):
    def __init__(self, size=256, random_crop=False, interpolation="bicubic"):
        super().__init__(data_csv="data/ade20k_examples.txt",
                         data_root="data/ade20k_images",
                         segmentation_root="data/ade20k_segmentations",
                         size=size, random_crop=random_crop,
                         interpolation=interpolation,
                         n_labels=151, shift_segmentation=False)


# With semantic map and scene label
class ADE20kBase(Dataset):
    def __init__(self, config=None, size=None, random_crop=False, interpolation="bicubic", crop_size=None):
        self.split = self.get_split()
        self.n_labels = 151 # unknown + 150
        self.data_csv = {"train": "data/ade20k_train.txt",
                         "validation": "data/ade20k_test.txt"}[self.split]
        self.data_root = "./data/ade20k_root"
        with open(os.path.join(self.data_root, "sceneCategories.txt"), "r") as f:
            self.scene_categories = f.read().splitlines()
        self.scene_categories = dict(line.split() for line in self.scene_categories)
        with open(self.data_csv, "r") as f:
            self.image_paths = f.read().splitlines()
        self._length = len(self.image_paths)
        ss = self.split
        if ss=='train':
            ss='training'
        self.labels = {
            "relative_file_path_": [l for l in self.image_paths],
            "file_path_": [os.path.join(self.data_root, "images",ss, l)
                           for l in self.image_paths],
            "relative_segmentation_path_": [l.replace(".jpg", ".png")
                                            for l in self.image_paths],
            "segmentation_path_": [os.path.join(self.data_root, "annotations",ss,
                                                l.replace(".jpg", ".png"))
                                   for l in self.image_paths],
            "scene_category": [self.scene_categories[l.replace(".jpg", "")]
                               for l in self.image_paths],
        }

        size = None if size is not None and size<=0 else size
        self.size = size
        if crop_size is None:
            self.crop_size = size if size is not None else None
        else:
            self.crop_size = crop_size
        if self.size is not None:
            self.interpolation = interpolation
            self.interpolation = {
                "nearest": cv2.INTER_NEAREST,
                "bilinear": cv2.INTER_LINEAR,
                "bicubic": cv2.INTER_CUBIC,
                "area": cv2.INTER_AREA,
                "lanczos": cv2.INTER_LANCZOS4}[self.interpolation]
            self.image_rescaler = albumentations.SmallestMaxSize(max_size=self.size,
                                                                 interpolation=self.interpolation)
            self.segmentation_rescaler = albumentations.SmallestMaxSize(max_size=self.size,
                                                                        interpolation=cv2.INTER_NEAREST)

        if crop_size is not None:
            self.center_crop = not random_crop
            if self.center_crop:
                self.cropper = albumentations.CenterCrop(height=self.crop_size, width=self.crop_size)
            else:
                self.cropper = albumentations.RandomCrop(height=self.crop_size, width=self.crop_size)
            self.preprocessor = self.cropper

    def __len__(self):
        return self._length

    def __getitem__(self, i):
        example = dict((k, self.labels[k][i]) for k in self.labels)
        image = Image.open(example["file_path_"])
        if not image.mode == "RGB":
            image = image.convert("RGB")
        image = np.array(image).astype(np.uint8)
        if self.size is not None:
            image = self.image_rescaler(image=image)["image"]
        segmentation = Image.open(example["segmentation_path_"])
        segmentation = np.array(segmentation).astype(np.uint8)
        if self.size is not None:
            segmentation = self.segmentation_rescaler(image=segmentation)["image"]
        if self.size is not None:
            processed = self.preprocessor(image=image, mask=segmentation)
        else:
            processed = {"image": image, "mask": segmentation}
        example["image"] = (processed["image"]/127.5 - 1.0).astype(np.float32)
        segmentation = processed["mask"]
        onehot = np.eye(self.n_labels)[segmentation]
        example["segmentation"] = onehot
        return example


class ADE20kTrain(ADE20kBase):
    # default to random_crop=True
    def __init__(self, config=None, size=None, random_crop=True, interpolation="bicubic", crop_size=None):
        super().__init__(config=config, size=size, random_crop=random_crop,
                          interpolation=interpolation, crop_size=crop_size)

    def get_split(self):
        return "train"


class ADE20kValidation(ADE20kBase):
    def get_split(self):
        return "validation"


if __name__ == "__main__":
    dset = ADE20kValidation()
    ex = dset[0]
    for k in ["image", "scene_category", "segmentation"]:
        print(type(ex[k]))
        try:
            print(ex[k].shape)
        except:
            print(ex[k])

[((83, 16, 83, 49), 'PIL.Image.open', 'Image.open', ({(83, 27, 83, 48): "example['file_path_']"}, {}), "(example['file_path_'])", False, 'from PIL import Image\n'), ((89, 23, 89, 64), 'PIL.Image.open', 'Image.open', ({(89, 34, 89, 63): "example['segmentation_path_']"}, {}), "(example['segmentation_path_'])", False, 'from PIL import Image\n'), ((65, 34, 66, 98), 'albumentations.SmallestMaxSize', 'albumentations.SmallestMaxSize', (), '', False, 'import albumentations\n'), ((67, 41, 68, 104), 'albumentations.SmallestMaxSize', 'albumentations.SmallestMaxSize', (), '', False, 'import albumentations\n'), ((99, 17, 99, 38), 'numpy.eye', 'np.eye', ({(99, 24, 99, 37): 'self.n_labels'}, {}), '(self.n_labels)', True, 'import numpy as np\n'), ((29, 18, 29, 69), 'os.path.join', 'os.path.join', ({(29, 31, 29, 45): 'self.data_root', (29, 47, 29, 68): '"""sceneCategories.txt"""'}, {}), "(self.data_root, 'sceneCategories.txt')", False, 'import os\n'), ((40, 27, 40, 71), 'os.path.join', 'os.path.join', ({(40, 40, 40, 54): 'self.data_root', (40, 56, 40, 64): '"""images"""', (40, 65, 40, 67): 'ss', (40, 69, 40, 70): 'l'}, {}), "(self.data_root, 'images', ss, l)", False, 'import os\n'), ((73, 31, 73, 101), 'albumentations.CenterCrop', 'albumentations.CenterCrop', (), '', False, 'import albumentations\n'), ((75, 31, 75, 101), 'albumentations.RandomCrop', 'albumentations.RandomCrop', (), '', False, 'import albumentations\n'), ((86, 16, 86, 31), 'numpy.array', 'np.array', ({(86, 25, 86, 30): 'image'}, {}), '(image)', True, 'import numpy as np\n'), ((90, 23, 90, 45), 'numpy.array', 'np.array', ({(90, 32, 90, 44): 'segmentation'}, {}), '(segmentation)', True, 'import numpy as np\n')]