# Copyright 2020 Amazon.com, Inc. or its affiliates. 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.
#   A copy of the License is located at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   or in the "license" file accompanying this file. This file 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 random
import numpy as np
import zipfile
import matplotlib.pyplot as plt

from PIL import Image
from mxnet.gluon.data.dataset import Dataset
import mxnet

from .data import MetaTaskDataContainer, TaskDataContainer
from .utils import list_dirs, list_files, random_affine_matrix, copy_dir, resize_imgs, download_url
from .config import DEFAULT_CONFIG_OMNIGLOT


class MetaTaskOmniglot(MetaTaskDataContainer):

    def __init__(self, root="./data", num_classes=None,
                 config=None, seed=1, context=None):

        """
        Onmiglot Data Container for Meta-Learning

        :param root: path containing the onmiglot dataset (if it does not exists, it
            download the data into this directory.
        :param num_classes: If not None, only those alphabet with at least num_classes
            are taken into consideration.
        :param config: If None, DEFAULT_CONFIG is loaded.
        :param seed: seed for random generator.
        """

        if config is None:
            self.config = DEFAULT_CONFIG_OMNIGLOT

        if context is None:
            context = mxnet.cpu()
        self.context = context

        self.url_dir = 'https://github.com/brendenlake/omniglot/raw/master/python/'
        self.urls = [
            os.path.join(self.url_dir, "images_background.zip"),
            os.path.join(self.url_dir, "images_evaluation.zip")
            ]

        self.root = root
        self.target_folder = "images_resized"

        self.download_dataset()

        self.seed = seed
        random.seed(self.seed)

        self.num_classes = num_classes

        self.alphabets = []
        self.alphabets_train = []
        self.alphabets_test = []
        self.alphabets_val = []

        self.target_path = os.path.join(self.root, self.target_folder)

        self._init_train_test_val_alphabets()

        hold = self.config["hold_out"]
        transform_image = self.config["transform_image"]
        transform_mxnet = self.config["transform_mxnet"]

        # Generate the training/test/val dataset.
        # Each dataset is a list of SubOmniglot objects (one per task)
        data_train = [TaskOmniglot(self.target_path, [a], num_classes, hold, transform_image=transform_image,
                                   transform_mxnet=transform_mxnet, context=context) for a in self.alphabets_train]
        data_test = [TaskOmniglot(self.target_path, [a], num_classes, hold, transform_image=transform_image,
                                  transform_mxnet=transform_mxnet, context=context) for a in self.alphabets_test]
        data_val = [TaskOmniglot(self.target_path, [a], num_classes, hold, transform_image=transform_image,
                                 transform_mxnet=transform_mxnet, context=context) for a in self.alphabets_val]

        super(MetaTaskOmniglot, self).__init__(data_train, data_test, data_val, context)

    def _init_train_test_val_alphabets(self):

        trs = self.config["num_tasks_train"]
        tes = self.config["num_tasks_test"]
        val = self.config["num_tasks_val"]

        # Filter out alphabet with less than num_classes, and shuffle
        self.alphabets = list(list_dirs(self.target_path))
        if self.num_classes:
            self.alphabets = []
            for a in self.alphabets:
                if len(list(list_dirs(os.path.join(self.target_path, a)))) >= self.num_classes:
                    self.alphabets.append(a)
            assert trs + tes < len(self.alphabets), 'cannot create test set'
        random.shuffle(self.alphabets)

        # init train/test/val alphabets
        self.alphabets_train = self.alphabets[:trs]
        self.alphabets_test = self.alphabets[trs:trs + tes]
        if val is not None:
            self.alphabets_val = self.alphabets[trs + tes:trs + tes + val]
        else:
            self.alphabets_val = self.alphabets[trs + tes:]

    def _copy_to_target(self, path_origin):
        out_path = os.path.abspath(os.path.join(self.root, self.target_folder))
        for alphabet in os.listdir(path_origin):
            copy_dir(os.path.join(path_origin, alphabet), os.path.join(out_path, alphabet))

    def _resize(self):
        resize_imgs(os.path.abspath(os.path.join(self.root, self.target_folder)))

    def download_dataset(self):
        """
        Helper function to download Onmiglot

        It download  both files, images_background.zip and images_evaluation, extract and extract
        them in self.target_path. It finally resizes all the image to 28X28 px.
        """
        if not os.path.exists(self.root):
            os.makedirs(self.root)

        for url in self.urls:
            file_name = url.split(os.sep)[-1]
            out_file = os.path.join(self.root, file_name)

            if not os.path.exists(out_file):
                out_file = download_url(url, out_file)
                with zipfile.ZipFile(out_file, "r") as zip_file:
                    zip_file.extractall(os.path.dirname(out_file))

                dir_images = os.path.abspath(out_file).split(".")[-2]
                self._copy_to_target(dir_images)
                self._resize()

    def plot_sample(self, num_samples, root="./sample_onmiglot"):

        """Plot N images from each alphabet and store the images in root."""

        if not os.path.exists(root):
            os.makedirs(root)

        fig_train = self._plot(num_samples, [dd._train_dataset for dd in self.train_tasks],
                               "Training Samples for Training Tasks")
        fig_train.savefig(os.path.join(root, "sample_train_train_tasks.png"))
        del fig_train
        fig_test = self._plot(num_samples, [dd._train_dataset for dd in self.test_tasks],
                              "Training Samples for Test Tasks")
        fig_test.savefig(os.path.join(root, "sample_train_test_tasks.png"))
        del fig_test
        fig_val = self._plot(num_samples, [dd._train_dataset for dd in self.val_tasks],
                             "Training Samples for Validation Tasks")
        fig_val.savefig(os.path.join(root, "sample_train_val_tasks.png"))
        del fig_val
        fig_train = self._plot(num_samples, [dd._val_dataset for dd in self.train_tasks],
                               "Validation Samples for Training Tasks")
        fig_train.savefig(os.path.join(root, "sample_val_train_tasks.png"))
        del fig_train
        fig_test = self._plot(num_samples, [dd._val_dataset for dd in self.test_tasks],
                              "Validation Samples for Test Tasks")
        fig_test.savefig(os.path.join(root, "sample_val_test_tasks.png"))
        del fig_test
        fig_val = self._plot(num_samples, [dd._val_dataset for dd in self.val_tasks],
                             "Validation Samples for Validation Tasks")
        fig_val.savefig(os.path.join(root, "sample_val_val_tasks.png"))
        del fig_val

    def _plot(self, num_samples, data, title):

        """Helper function for plotting."""

        num_alphabets = len(data)
        fig, ax = plt.subplots(num_alphabets, num_samples, figsize=(2 * num_samples, 2 * num_alphabets))
        fig.subplots_adjust(left=None, bottom=None, right=None, top=None, wspace=0.5, hspace=0.5)
        for mm in range(num_alphabets):
            ix_perm = np.random.permutation(len(data[mm]))
            for nn in range(num_samples):
                im, label = data[mm][ix_perm[nn]]
                ax[mm, nn].imshow(im.asnumpy())
                ax[mm, nn].set_title("Label: {} \n ({})".format(label, data[mm].items[nn][0].split("/")[-3][0:5]))
                ax[mm, nn].axis("off")
        fig.suptitle(title, size=18)
        return fig


class TaskOmniglot(TaskDataContainer):
    """
    Onmiglot Task Container

    Represent a single Onmiglot Task (can contain one or more alphabets)
    """

    def __init__(self, path, alphabets, num_classes=None, hold_out=None,
                 transform_image=None, transform_mxnet=None, seed=None, context=None):
        """

        :param path: Root path
        :param alphabets: Alphabet to be considered for the task
        :param num_classes: Number of classes/characters per alphabet
        :param hold_out: Number of images to hold out for validation
        :param transform_image: Applied to the PIL image (before loading it in an mx.ndarray)
        :param transform_mxnet: Applied to the mx.ndarray when loading a character
        :param seed: seed for the random generator
        """

        if context is None:
            context = mxnet.cpu()
        self.context = context

        self.path = path
        self.alphabets = alphabets
        self.num_classes = num_classes
        self.hold_out = hold_out
        self.transform_image = transform_image
        self.transform_mxnet = transform_mxnet
        self.target_transform = None
        self.seed = seed

        # Generate a list of alphabets and characters
        self._alphabets = [a for a in list_dirs(self.path) if a in self.alphabets]
        self._characters = sum([[os.path.join(a, c) for c in list_dirs(os.path.join(self.path, a))]
                                for a in self._alphabets], [])

        if seed:
            random.seed(seed)

        random.shuffle(self._characters)

        if self.num_classes:
            self._characters = self._characters[:num_classes]

        self._train_character_images = []
        self._val_character_images = []
        for idx, character in enumerate(self._characters):
            train_characters = []
            val_characters = []

            for img_count, image in enumerate(list_files(os.path.join(self.path, character), 'png')):
                if hold_out and img_count < hold_out:
                    val_characters.append((image, idx))
                else:
                    train_characters.append((image, idx))
            self._train_character_images.append(train_characters)
            self._val_character_images.append(val_characters)

        self._flat_train_character_images = sum(self._train_character_images, [])
        self._flat_val_character_images = sum(self._val_character_images, [])

        train_dataset = ImageListDataset(self._flat_train_character_images,
                                         flag=0, transform_image=transform_image,
                                         transform_mxnet=transform_mxnet, context=context)

        if len(self._flat_val_character_images) > 0:
            val_dataset = ImageListDataset(self._flat_val_character_images, flag=0,
                                           transform_image=transform_image,
                                           transform_mxnet=transform_mxnet, context=context)
        else:
            val_dataset = None

        super(TaskOmniglot, self).__init__(train_dataset, val_dataset)


class ImageListDataset(Dataset):

    def __init__(self, items, flag=1, transform_image=None, transform_mxnet=None, context=None):
        self.items = items
        self._flag = flag
        self._transform_image = transform_image
        self._transform_mxnet = transform_mxnet

        if context is None:
            context = mxnet.cpu()
        self.context = context

    def __getitem__(self, idx):
        img = Image.open(self.items[idx][0], mode='r').convert('L')
        if self._transform_image is not None:
            img = random_affine_matrix(img)
        img = mxnet.nd.array(img, ctx=self.context)
        label = self.items[idx][1]
        if self._transform_mxnet is not None:
            return self._transform_mxnet(img, label)
        return img, label

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


if __name__ == '__main__':

    o = MetaTaskOmniglot()
    o.plot_sample(5)

    meta_batch_size = 2
    batch_size = 20
    train_tasks = o.train_tasks

    assert len(train_tasks) == 3

    for task in train_tasks:
        tr_iterator = task.get_train_iterator(batch_size)
        for data in tr_iterator:
            assert (data[0].shape == (batch_size, 28, 28))
            assert (data[1].shape == (batch_size, ))
            assert (data[1].asnumpy().dtype == np.int)
            break

        val_iterator = task.get_val_iterator(batch_size)
        for data in val_iterator:
            assert (data[0].shape == (batch_size, 28, 28))
            assert (data[1].shape == (batch_size, ))
            assert (data[1].asnumpy().dtype == np.int)
            break
