#!/usr/bin/env python3

import torch
import os
from typing import List
from advattack.error_handling.exception import DatasetNotFoundError
from advattack.data_handling.dataset import Dataset
import matplotlib.pyplot as plt
import glob
from advattack.util.logger import logger
import numpy as np
import pickle

class Cifar10Dataset(Dataset):
    """ Cifar dataset (https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz)
    """

    # TODO: Fix this ugly hard coding ...
    logger = logger.getChild("data_handling.dataset.cifar.cifar_dataset.CifarDataset")

    def __init__(self, root_path, feature_transform_fun=None, target_transform_fun=None):
        super(Cifar10Dataset, self).__init__(root_path, feature_transform_fun, target_transform_fun)

    def __len__(self):
        return self.samples.shape[0]

    def visualize_samples(self, min_index, max_index):
        plots_per_column = 15 # we might want to make this configurable
        plot_count = max_index - min_index + 1
        cols = min(plot_count, plots_per_column)
        rows = int(plot_count/plots_per_column) + 1
        plt.figure(figsize=(plots_per_column, rows))
        plt.subplots_adjust(top=0.8, bottom=0, hspace=1, wspace=0.5)
        for fig_index, index in enumerate(range(min_index, max_index+1)):
            ax = plt.subplot(rows, cols, fig_index+1)
            ax.set_axis_off()
            pixels, label_index = self[index]
            label = self.label_map[label_index]
            plt.title(f'{index}\n{label}')
            plt.imshow(pixels)
        plt.show()

    @classmethod
    def get_config(cls):
        # {resource_url: (md5_hash, destination_folder, extraction_function)}
        resources = {
            'https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz': ("c58f30108f718f92721af3b95e74349a", "", Cifar10Dataset.extract_dataset)
        }
        return resources

    @classmethod
    def load(cls, path, feature_transform_fun=None, target_transform_fun=None):
        if cls.check_exists(path):
            return Cifar10Dataset(root_path=path)
        else:
            raise DatasetNotFoundError(f"Dataset not found in {path}.")

    @classmethod
    def check_exists(cls, root_path) -> bool:
        batch_files = glob.glob(os.path.join(root_path, "**/data_batch_*"))
        label_map_files = glob.glob(os.path.join(root_path, "**/batches.meta"))

        return (len(batch_files) == 5) and (len(label_map_files) == 1)

    @classmethod
    def create_dataset(cls, root_path) -> str:
        resource_extraction_mapping = cls.get_config()
        cls.populate_data_repository(root_path, resource_extraction_mapping, force_download=True)
        return root_path

    @classmethod
    def extract_dataset(cls, source_path):
        unzipped_path = Dataset.extract_tar(source_path, remove_finished=True)

    def load_data_from_disc(self) -> (List, List, List):
        """Method that implements loading functionality of an on disk dataset.
        :return: samples, targets
        """
        batch_files = glob.glob(os.path.join(self.root_path, "**/*_batch*"))
        samples = []
        labels = []
        Cifar10Dataset.logger.debug(f"Loading samples from {self.root_path}")
        for batch_file in batch_files:
            batch_samples, batch_labels = self.load_image_batch(batch_file)
            samples.append(batch_samples)
            labels.extend(batch_labels)
        samples = np.vstack(samples).reshape(-1, 3, 32, 32)
        samples = samples.transpose((0, 2, 3, 1))  # convert from CHW to HWC format (Height Width Channel)
        samples = torch.tensor(samples).float()

        # load labels
        label_map_path = glob.glob(os.path.join(self.root_path, "**/batches.meta"))[0]
        label_map = self.load_label_map(label_map_path)
        return samples, labels, label_map

    def load_image_batch(self, batch_path):
        with open(batch_path, 'rb') as f:
            entry = pickle.load(f, encoding='latin1')
        targets = entry['labels']
        samples = entry["data"]
        return samples, targets

    def load_label_map(self, label_map_path):
        with open(label_map_path, 'rb') as f:
            data = pickle.load(f, encoding='latin1')
        label_map = data["label_names"]
        return label_map


if __name__ == "__main__":
    from advattack import datasets_path
    root_path = os.path.join(datasets_path, Cifar10Dataset.get_dataset_identifier())
    if not Cifar10Dataset.check_exists(root_path):
        root_path = Cifar10Dataset.create_dataset(root_path=root_path)

    dataset = Cifar10Dataset.load(root_path)
    dataset.visualize_samples(min_index=0, max_index=50)
    len(dataset)


