# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import copy

import cv2
import numpy as np

from paddleseg.cvlibs import manager
from paddleseg.transforms import Compose
from paddleseg.datasets import Dataset
from paddleseg.utils.download import download_file_and_uncompress
from paddleseg.utils import seg_env
import paddleseg.transforms.functional as F

URL = "https://paddleseg.bj.bcebos.com/dataset/Supervisely_face.zip"


@manager.DATASETS.add_component
class SUPERVISELY(Dataset):
    """
    Supervise.ly dataset `https://supervise.ly/`.

    Args:
        common_transforms (list): A list of common image transformations for two inputs of portrait net.
        transforms1 (list): A list of image transformations for the first input of portrait net.
        transforms2 (list): A list of image transformations for the second input of portrait net.
        dataset_root (str, optional): The Supervise.ly dataset directory. Default: None.
        mode (str, optional): A subset of the entire dataset. It should be one of ('train', 'val'). Default: 'train'.
        edge (bool, optional): Whether to compute edge while training. Default: False
    """
    NUM_CLASSES = 2

    def __init__(self,
                 common_transforms,
                 transforms1,
                 transforms2,
                 dataset_root=None,
                 mode='train',
                 edge=False):
        self.dataset_root = dataset_root
        self.common_transforms = Compose(common_transforms)
        self.transforms = self.common_transforms
        if transforms1 is not None:
            self.transforms1 = Compose(transforms1, to_rgb=False)
        if transforms2 is not None:
            self.transforms2 = Compose(transforms2, to_rgb=False)
        mode = mode.lower()
        self.ignore_index = 255
        self.mode = mode
        self.num_classes = self.NUM_CLASSES
        self.input_width = 224
        self.input_height = 224

        if self.dataset_root is None:
            self.dataset_root = download_file_and_uncompress(
                url=URL,
                savepath=seg_env.DATA_HOME,
                extrapath=seg_env.DATA_HOME)
        elif not os.path.exists(self.dataset_root):
            self.dataset_root = os.path.normpath(self.dataset_root)
            savepath, extraname = self.dataset_root.rsplit(
                sep=os.path.sep, maxsplit=1)
            self.dataset_root = download_file_and_uncompress(
                url=URL,
                savepath=savepath,
                extrapath=savepath,
                extraname=extraname)

        if mode == 'train':
            path = os.path.join(dataset_root, 'supervisely_face_train_easy.txt')
        else:
            path = os.path.join(dataset_root, 'supervisely_face_test_easy.txt')
        with open(path, 'r') as f:
            files = f.readlines()
        files = ["/".join(file.split('/')[1:]) for file in files]
        img_files = [os.path.join(dataset_root, file).strip() for file in files]
        label_files = [
            os.path.join(dataset_root, file.replace('/img/', '/ann/')).strip()
            for file in files
        ]

        self.file_list = [
            [img_path, label_path]
            for img_path, label_path in zip(img_files, label_files)
        ]

    def __getitem__(self, item):
        image_path, label_path = self.file_list[item]
        im = cv2.imread(image_path)
        label = cv2.imread(label_path, 0)
        label[label > 0] = 1

        if self.mode == "val":
            common_im, label = self.common_transforms(im=im, label=label)
            im = np.float32(common_im[::-1, :, :])  # RGB => BGR
            im_aug = copy.deepcopy(im)
        else:
            common_im, label = self.common_transforms(im=im, label=label)
            common_im = np.transpose(common_im, [1, 2, 0])
            # add augmentation
            im, _ = self.transforms1(common_im)
            im_aug, _ = self.transforms2(common_im)

            im = np.float32(im[::-1, :, :])  # RGB => BGR
            im_aug = np.float32(im_aug[::-1, :, :])  # RGB => BGR

        label = cv2.resize(
            np.uint8(label), (self.input_width, self.input_height),
            interpolation=cv2.INTER_NEAREST)

        # add mask blur
        label = np.uint8(cv2.blur(label, (5, 5)))
        label[label >= 0.5] = 1
        label[label < 0.5] = 0

        edge_mask = F.mask_to_binary_edge(
            label, radius=4, num_classes=self.num_classes)
        edge_mask = np.transpose(edge_mask, [1, 2, 0]).squeeze(axis=-1)
        im = np.concatenate([im_aug, im])
        if self.mode == "train":
            return im, label, edge_mask
        else:
            return im, label