# 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 glob
import random

import paddle
import numpy as np
from PIL import Image

from paddleseg.cvlibs import manager
from paddleseg.transforms import Compose

# Random seed is set to ensure that after shuffling dataset per epoch during multi-gpu training, the data sequences of all gpus are consistent.
random.seed(100)


@manager.DATASETS.add_component
class AutoNueCrop(paddle.io.Dataset):
    """
    You need to to download and convert the [India Driving Dataset](https://idd.insaan.iiit.ac.in/evaluation/autonue21/#bm5)
    following the instructions of Segmentation Track. IDD_Dectection dataset also need for pseudo-labeling.

    And then, you need to organize data following the below structure.

    IDD_Segmentation
    |
    |--leftImg8bit
    |  |--train
    |  |--val
    |  |--test
    |
    |--gtFine
    |  |--train
    |  |--val
    |  |--test

    Args:
        transforms (list): Transforms for image.
        dataset_root (str): Cityscapes dataset directory.
        mode (str, optional): Which part of dataset to use. it is one of ('train', 'val', 'test'). Default: 'train'.
        coarse_multiple (float|int, optional): Multiple of the amount of coarse data relative to fine data. Default: 1
        add_val (bool, optional): Whether to add val set in training. Default: False
    """

    def __init__(self,
                 transforms,
                 dataset_root,
                 mode='train',
                 coarse_multiple=1,
                 add_val=False):
        self.dataset_root = dataset_root
        self.transforms = Compose(transforms)
        self.file_list = list()
        mode = mode.lower()
        self.mode = mode
        self.num_classes = 26
        self.ignore_index = 255
        self.coarse_multiple = coarse_multiple

        if mode not in ['train', 'val', 'test']:
            raise ValueError(
                "mode should be 'train', 'val' or 'test', but got {}.".format(
                    mode))

        if self.transforms is None:
            raise ValueError("`transforms` is necessary, but it is None.")

        img_dir = os.path.join(self.dataset_root, 'leftImg8bit')
        label_dir = os.path.join(self.dataset_root, 'gtFine')
        if self.dataset_root is None or not os.path.isdir(
                self.dataset_root) or not os.path.isdir(
                    img_dir) or not os.path.isdir(label_dir):
            raise ValueError(
                "The dataset is not Found or the folder structure is nonconfoumance."
            )

        label_files = sorted(
            glob.glob(
                os.path.join(label_dir, mode, '*',
                             '*_gtFine_labellevel3Ids.png')))
        img_files = sorted(
            glob.glob(os.path.join(img_dir, mode, '*', '*_leftImg8bit.*')))

        self.file_list = [
            [img_path, label_path]
            for img_path, label_path in zip(img_files, label_files)
        ]
        # for ii in range(len(self.file_list)):
        #     print(self.file_list[ii])
        # print(len(self.file_list))
        self.num_files = len(self.file_list)
        self.total_num_files = self.num_files

        if mode == 'train':
            # whether to add val set in training
            if add_val:
                label_files = sorted(
                    glob.glob(
                        os.path.join(label_dir, 'val', '*',
                                     '*_gtFine_labellevel3Ids.png')))
                img_files = sorted(
                    glob.glob(
                        os.path.join(img_dir, 'val', '*', '*_leftImg8bit.*')))
                val_file_list = [
                    [img_path, label_path]
                    for img_path, label_path in zip(img_files, label_files)
                ]
                self.file_list.extend(val_file_list)
                for ii in range(len(self.file_list)):
                    print(self.file_list[ii])
                print(len(self.file_list))
                self.num_files = len(self.file_list)
                self.total_num_files = self.num_files

    def __getitem__(self, idx):
        if self.mode == 'test':
            image_path, label_path = self.file_list[idx]
            im, _ = self.transforms(im=image_path)
            (h, w) = im.shape[1:]
            im1 = im[:, 540:540 + 720, 320:320 + 1280]
            im2 = im[:, 540:540 + 720, 960:960 + 1280]
            im3 = im[:, 540:540 + 720, 1600:1600 + 1280]
            return im1, im2, im3, image_path
        elif self.mode == 'val':
            image_path, label_path = self.file_list[idx]
            im, label = self.transforms(im=image_path, label=label_path)
            (h, w) = im.shape[1:]
            im1 = im[:, 540:540 + 720, 320:320 + 1280]
            im2 = im[:, 540:540 + 720, 960:960 + 1280]
            im3 = im[:, 540:540 + 720, 1600:1600 + 1280]
            label = label[540:540 + 720, 320:1600 + 1280]
            return im1, im2, im3, label
        else:
            if idx >= self.num_files:
                image_path, label_path = self.coarse_file_list[idx -
                                                               self.num_files]
            else:
                image_path, label_path = self.file_list[idx]

            im, label = self.transforms(im=image_path, label=label_path)
            (h, w) = im.shape[1:]

            start_w = np.linspace(320, 1600, 5).tolist()
            np.random.shuffle(start_w)
            start = int(start_w[0])
            crop_im = im[:, 540:540 + 720, start:(start + 1280)]
            crop_label = label[540:540 + 720, start:(start + 1280)]
            return crop_im, crop_label

    def __len__(self):
        return self.total_num_files