# Copyright (C) 2020-2021 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

"""You may copy this file as the starting point of your own model."""

from collections.abc import Iterable
from logging import getLogger
from os import makedirs
from pathlib import Path
from urllib.request import urlretrieve
from zipfile import ZipFile

import numpy as np
import torch
from torch.utils.data import random_split
from torchvision.datasets import ImageFolder
from torchvision.transforms import ToTensor
from tqdm import tqdm

from openfl.federated import PyTorchDataLoader
from openfl.utilities import validate_file_hash

logger = getLogger(__name__)


class PyTorchHistologyInMemory(PyTorchDataLoader):
    """PyTorch data loader for Histology dataset."""

    def __init__(self, data_path=None, batch_size=32, **kwargs):
        """Instantiate the data object.

        Args:
            data_path: The file path to the data. If None, initialize for model creation only.
            batch_size: The batch size of the data loader
            **kwargs: Additional arguments, passed to super init
             and load_mnist_shard
        """
        super().__init__(batch_size, random_seed=0, **kwargs)

        # Set Histology-specific default attributes
        self.feature_shape = [3, 150, 150]
        self.num_classes = 8

        # If data_path is None, this is being used for model initialization only
        if data_path is None:
            logger.info("Initializing dataloader for model creation only (no data loading)")
            return

        try:
            int(data_path)
        except ValueError:
            raise ValueError(
                "Expected `%s` to be representable as `int`, as it refers to the data shard " +
                "number used by the collaborator.",
                data_path
            )

        X_train, y_train, X_valid, y_valid = load_histology_shard(
            shard_num=int(data_path),
            feature_shape=self.feature_shape,
            num_classes=self.num_classes,
            **kwargs
        )
        self.X_train = X_train
        self.y_train = y_train
        self.X_valid = X_valid
        self.y_valid = y_valid

    def get_feature_shape(self):
        """Returns the shape of an example feature array.

        Returns:
            list: The shape of an example feature array [3, 150, 150] for Histology images.
        """
        return self.feature_shape

    def get_num_classes(self):
        """Returns the number of classes for classification tasks.

        Returns:
            int: The number of classes (8 for Histology dataset).
        """
        return self.num_classes


class HistologyDataset(ImageFolder):
    """Colorectal Histology Dataset."""

    URL = ('https://zenodo.org/record/53169/files/Kather_'
           'texture_2016_image_tiles_5000.zip?download=1')
    FILENAME = 'Kather_texture_2016_image_tiles_5000.zip'
    FOLDER_NAME = 'Kather_texture_2016_image_tiles_5000'
    ZIP_SHA384 = ('7d86abe1d04e68b77c055820c2a4c582a1d25d2983e38ab724e'
                  'ac75affce8b7cb2cbf5ba68848dcfd9d84005d87d6790')
    DEFAULT_PATH = Path.cwd().absolute() / 'data'

    def __init__(self, root: Path = DEFAULT_PATH, **kwargs) -> None:
        """Initialize."""
        makedirs(root, exist_ok=True)
        filepath = root / HistologyDataset.FILENAME
        if not filepath.is_file():
            self.pbar = tqdm(total=None)
            urlretrieve(HistologyDataset.URL, filepath, self.report_hook)  # nosec
            validate_file_hash(filepath, HistologyDataset.ZIP_SHA384)
            with ZipFile(filepath, 'r') as f:
                f.extractall(root)

        super(HistologyDataset, self).__init__(root / HistologyDataset.FOLDER_NAME, **kwargs)

    def report_hook(self, count, block_size, total_size):
        """Update progressbar."""
        if self.pbar.total is None and total_size:
            self.pbar.total = total_size
        progress_bytes = count * block_size
        self.pbar.update(progress_bytes - self.pbar.n)

    def __getitem__(self, index):
        """Allow getting items by slice index."""
        if isinstance(index, Iterable):
            return [super(HistologyDataset, self).__getitem__(i) for i in index]
        else:
            return super(HistologyDataset, self).__getitem__(index)


def one_hot(labels, classes):
    """
    One Hot encode a vector.

    Args:
        labels (list):  List of labels to onehot encode
        classes (int): Total number of categorical classes

    Returns:
        np.array: Matrix of one-hot encoded labels
    """
    return np.eye(classes)[labels]


def _load_raw_datashards(shard_num, collaborator_count, train_split_ratio=0.8):
    """
    Load the raw data by shard.

    Returns tuples of the dataset shard divided into training and validation.

    Args:
        shard_num (int): The shard number to use
        collaborator_count (int): The number of collaborators in the federation

    Returns:
        2 tuples: (image, label) of the training, validation dataset
    """
    dataset = HistologyDataset(transform=ToTensor())
    n_train = int(train_split_ratio * len(dataset))
    n_valid = len(dataset) - n_train
    ds_train, ds_val = random_split(
        dataset, lengths=[n_train, n_valid], generator=torch.manual_seed(0))

    # create the shards
    X_train, y_train = list(zip(*ds_train[shard_num::collaborator_count]))
    X_train, y_train = np.stack(X_train), np.array(y_train)

    X_valid, y_valid = list(zip(*ds_val[shard_num::collaborator_count]))
    X_valid, y_valid = np.stack(X_valid), np.array(y_valid)

    return (X_train, y_train), (X_valid, y_valid)


def load_histology_shard(shard_num, collaborator_count, feature_shape=None, num_classes=None,
                         categorical=False, channels_last=False, **kwargs):
    """
    Load the Histology dataset.

    Args:
        shard_num (int): The shard to use from the dataset
        collaborator_count (int): The number of collaborators in the federation
        feature_shape (list, optional): The shape of input features.
        num_classes (int, optional): Number of classes.
        categorical (bool): True = convert the labels to one-hot encoded
         vectors (Default = True)
        channels_last (bool): True = The input images have the channels
         last (Default = True)
        **kwargs: Additional parameters to pass to the function

    Returns:
        numpy.ndarray: The training data
        numpy.ndarray: The training labels
        numpy.ndarray: The validation data
        numpy.ndarray: The validation labels
    """
    img_rows, img_cols = feature_shape[1], feature_shape[2]

    (X_train, y_train), (X_valid, y_valid) = _load_raw_datashards(
        shard_num, collaborator_count)

    if channels_last:
        X_train = X_train.reshape(X_train.shape[0], img_rows, img_cols, 3)
        X_valid = X_valid.reshape(X_valid.shape[0], img_rows, img_cols, 3)
    else:
        X_train = X_train.reshape(X_train.shape[0], 3, img_rows, img_cols)
        X_valid = X_valid.reshape(X_valid.shape[0], 3, img_rows, img_cols)

    logger.info(f'Histology > X_train Shape : {X_train.shape}')
    logger.info(f'Histology > y_train Shape : {y_train.shape}')
    logger.info(f'Histology > Train Samples : {X_train.shape[0]}')
    logger.info(f'Histology > Valid Samples : {X_valid.shape[0]}')

    if categorical:
        # convert class vectors to binary class matrices
        y_train = one_hot(y_train, num_classes)
        y_valid = one_hot(y_valid, num_classes)

    return X_train, y_train, X_valid, y_valid
