"""
IBM Confidential
OCO Source Materials
5737-H76, 5725-W78, 5900-A1R
(c) Copyright IBM Corp. 2020 All Rights Reserved.
The source code for this program is not published or otherwise divested of its trade secrets,
irrespective of what has been deposited with the U.S. Copyright Office.
"""
import torch
import logging
import time
import copy
from importlib import import_module

import numpy as np
import torch.onnx
from sklearn import metrics
from skorch import NeuralNet
from skorch.callbacks import EpochScoring
from skorch.exceptions import NotInitializedError

from ibmfl.model.fl_model import FLModel
from ibmfl.model.model_update import ModelUpdate
from ibmfl.util import config
from ibmfl.exceptions import LocalTrainingException, ModelException, FLException

logger = logging.getLogger(__name__)


class PytorchFLModel(FLModel):
    """
    Wrapper class for importing a Pytorch based model
    """

    def __init__(self, model_name,
                 model_spec=None,
                 pytorch_module=None,
                 module_init_params=None,
                 **kwargs):
        """
        Create a `PytorchFLModel` instance from a Pytorch model.
        If pytorch_model is provided, it will use it; otherwise it will take
        the model_spec to create the model.

        :param model_name: String specifying the type of model e.g., Pytorch_NN
        :type model_name: `str`
        :param model_spec: A dictionary specifying path to saved nn.sequence container
        :type model_spec: 'dict'
        :param pytorch_module: uninstantiated pytorch model class
        :type pytorch_module: torch.nn.Module class reference \
        :param module_init_params: A dictionary with the values for the model's init \
        arguments. The key for the init parameters must be prefixed with module__ \
        (ex. if the parameter name is hidden_size, then key must be module__hidden_size)
        :type module_init_params: 'dict'
        :param kwargs: A dictionary contains other parameter settings on \
         to initialize a PyTorch model.
        :type kwargs: `dict`
        """
        super().__init__(model_name, model_spec, **kwargs)

        self.model_type = 'PyTorch'
        self.model = None
        self.module = None
        self.optimizer = torch.optim.SGD
        self.criterion = torch.nn.NLLLoss
        if pytorch_module is None:
            if model_spec is None or (not isinstance(model_spec, dict)):
                raise ValueError('Initializing model requires '
                                 'a model specification or uninstantiated '
                                 'pytorch model class reference '
                                 'None was provided')
            # In this case we need to recreate the model from model_spec
            self.module = self.load_model_from_spec(model_spec)
        else:
            self.module = pytorch_module
        self.model = self.initialize_model(self.module, self.optimizer,
                                           self.criterion, module_init_params)

        if self.use_gpu_for_training and torch.cuda.device_count() > 0:
            if self.num_gpus > torch.cuda.device_count():
                logger.error('Selected number of gpus to use for training exceeds number of available gpus, ' +
                             str(torch.cuda.device_count()) +
                             'Set number of gpus to maximum available on device.')
                self.num_gpus = torch.cuda.device_count()

            device_ids = list(range(self.num_gpus))
            self.model.module_ = torch.nn.DataParallel(self.model.module_, device_ids=device_ids)
            self.model.set_params(device='cuda')
            self.model.module_.to('cuda')

    def initialize_model(self, pytorch_module, optimizer, criterion,
                         module_init_params=None):
        """
        Initializes a pytorch model via skorch library

        :param pytorch_module: uninstantiated pytorch model class
        :type pytorch_module: torch.nn.Module class reference
        :param optimizer: the optimizer to use
        :type optimizer: pytorch optimizer class
        :param criterion: the loss function to use
        :type criterion: pytorch loss function class
        :param module_init_params: A dictionary with the values for the model's init arguments. \
        The key for the init parameters must be prefixed with module__  \
        (ex. if the parameter name is hidden_size, then key must be module__hidden_size)
        :type module_init_params: 'dict'
        :return: an initialized skorch model
        :rtype: 'skorch.NeuralNet'
        """
        if module_init_params is None:
            module_init_params = {}
        model = NeuralNet(
            module=pytorch_module,
            optimizer=optimizer,
            criterion=criterion,
            warm_start=True,
            callbacks=[('valid_acc', EpochScoring(self.valid_acc,
                                                  lower_is_better=False,
                                                  use_caching=True,
                                                  name='test_valid_acc'))],
            **module_init_params,
        )
        model.initialize()
        return model

    def fit_model(self, train_data, fit_params=None, **kwargs):
        """
        Fits current model with provided training data.

        :param train_data: Training data, a tuple \
        given in the form (x_train, y_train). otherwise, input compatible with skorch.dataset.Dataset
        :type train_data: `(np.ndarray, ”np.ndarray)`
        :param fit_params: (optional) Dictionary with hyperparameters \
        that will be used to call fit function. \
        Hyperparameter parameters should match pytorch expected values \
        e.g., `epochs`, which specifies the number of epochs to be run. \
        If no `epochs` or `batch_size` are provided, a default value \
        will be used (1 and 128, respectively).
        :type fit_params: `dict`
        :return: None
        """
        # Initialized with default values
        batch_size = 128
        epochs = 1
        lr = 0.01
        optimizer_params = {}
        hyperparams = fit_params.get('hyperparams', {}) if fit_params else {}

        if hyperparams:
            local_params = hyperparams.get('local', {}) or {}
            training_hp = local_params.get('training', {}) or {}
            
            epochs = training_hp.get('epochs', epochs)
            batch_size = training_hp.get('batch_size', batch_size)
            lr = training_hp.get('lr', lr)
            
            if 'optimizer' in local_params:
                optimizer = local_params['optimizer']
                if isinstance(optimizer, str):
                    try:
                        optimizer = getattr(import_module(
                            'torch.optim'), optimizer.split('.')[-1])

                    except Exception as e:
                        optimizer = self.optimizer
                        logger.exception(str(e))
                        print('selected optimizer not found. using default')
                if optimizer != self.optimizer:
                    self.optimizer = optimizer
                    self.model.set_params(optimizer=self.optimizer)

                    if self.use_gpu_for_training and torch.cuda.device_count() > 0:
                        self.model.module_.to(self.model.device)

            if 'criterion' in local_params:
                criterion = local_params['criterion']
                if isinstance(criterion, str):
                    try:
                        criterion = getattr(import_module(
                            'torch.nn'), criterion.split('.')[-1])

                    except Exception as e:
                        criterion = self.criterion
                        logger.exception(str(e))
                        print('selected criterion not found. using default')

                if criterion != self.criterion:
                    self.criterion = criterion
                    self.model.set_params(criterion=self.criterion)

            if 'optimizer_params' in local_params:
                optimizer_params = local_params['optimizer_params']
        self.model.set_params(batch_size=batch_size, lr=lr, **optimizer_params)

        try:
            if type(train_data) is tuple:
                # Extract x_train and y_train, by default,
                # label is stored in the last column
                x = train_data[0]
                y = train_data[1]
                self.model.fit_loop(x, y, epochs=epochs)

            else:
                # otherwise, expect that input is a pytorch Dataset generator
                self.model.fit_loop(train_data, epochs=epochs)

        except Exception as e:
            logger.exception(str(e))
            raise LocalTrainingException(
                'Error occurred while performing model.fit')

    def update_model(self, model_update):
        """
        Update model with provided model_update, where model_update
        should be generated according to `PytorchFLModel.get_model_update()`.

        :param model_update: `ModelUpdate` object that contains the weights \
        that will be used to update the model.
        :type model_update: `ModelUpdate`
        :return: None
        """
        if isinstance(model_update, ModelUpdate):
            for p1, p2 in zip(self.get_weights(), model_update.get('weights')):
                p1.data = torch.from_numpy(p2)
                p1.data.requires_grad = True

            if self.use_gpu_for_training and torch.cuda.device_count() > 0:
                self.model.module_.to(self.model.device)
        else:
            raise ValueError('Provided model_update should be of type Model.'
                             'Instead they are:{0}'.format(str(type(model_update))))

    def get_model_update(self):
        """
        Generates a `ModelUpdate` object that will be sent to other entities.

        :return: ModelUpdate
        :rtype: `ModelUpdate`
        """
        weights = self.get_weights(to_numpy=True)
        update = {'weights': weights}

        return ModelUpdate(**update)

    def predict(self, x):
        """
        Perform prediction for a batch of inputs. Note that for classification
        problems, it returns the resulting probabilities.

        :param x: Samples with shape as expected by the model.
        :type x: `np.ndarray`, input compatible with skorch.dataset.Dataset, or pytorch dataloader
        :return: Array of predictions
        :rtype: `np.ndarray`
        """
        if isinstance(x, torch.utils.data.DataLoader):
            return self.predict_generator(x)

        return self.model.predict(x)

    def predict_generator(self, dataloader):
        """
        Performs predictions using a pytorch dataloader

        :param x: A pytorch dataloader with a dataset for predicting
        :type x: 'torch.utils.data.DataLoader'
        :return: Array of predictions
        :rtype: 'np.ndarray'
        """
        predictions = None
        for x_batch, _ in dataloader:
            y_pred = self.model.evaluation_step(x_batch)
            y_pred = y_pred.numpy()
            if predictions is None:
                predictions = y_pred
            else:
                predictions = np.append(predictions, y_pred, 0)

        return predictions

    def evaluate(self, test_dataset, eval_metrics=None, **kwargs):
        """
        Evaluates the model given testing data.
        :param test_dataset: Testing data, a tuple given in the form \
        (x_test, y_test) or a pytorch DataLoader
        :type test_dataset: `np.ndarray`
        :param eval_metrics: a list of sklearn.metric class, or a function for evaluating a pytorch dataloader batch
        :type eval_metrics: 'sklearn.metrics' or 'function'
        :param kwargs: Dictionary of metrics available for the model
        :type kwargs: `dict`
        """

        if type(test_dataset) is tuple:
            x_test = test_dataset[0]
            y_test = test_dataset[1]

            return self.evaluate_model(x_test, y_test, eval_metrics)

        else:
            return self.evaluate_generator_model(test_dataset, eval_metrics)

    def evaluate_model(self, x, y, eval_metrics=None, **kwargs):
        """
        Evaluates the model given x and y.

        :param x:  Samples with shape as expected by the model.
        :type x: `np.ndarray`
        :param y: Corresponding labels to x
        :type y: `np.ndarray`
        :param eval_metrics: A list of sklearn.metric class or a functions with the signature \
                        (y_true, y_pred, **kwargs) => float
        :type eval_metrics: `sklearn.metrics` or 'function'
        :return: dictionary of metrics
        :rtype: 'dict'
        """
        if eval_metrics is None or len(eval_metrics) == 0:
            eval_metrics = [metrics.accuracy_score]

        # get names for eval_metrics
        metric_names = [getattr(metric, '__name__', repr(metric))
                        for metric in eval_metrics]

        y_pred = self.predict(x)
        y_pred_exp = np.exp(y_pred)
        y_pred_argmax = np.argmax(y_pred_exp, axis=1)

        try:
            metric_dict = dict(zip(metric_names, eval_metrics))
            # NOTE: Had to replace comprehension with for loop to tackle loss fn case which uses y_pred not argmax
            for metric, fn in metric_dict.items():
                if metric is 'loss':
                    metric_dict[metric] = fn(y, y_pred, self.model)
                else:
                    metric_dict[metric] = fn(y, y_pred_argmax, **kwargs)
            # metric_dict = {metric_name: metric(y, y_pred, **kwargs) for metric_name, metric in
            #                zip(metric_names, eval_metrics)}
            return metric_dict
        except TypeError as exc:
            logger.exception(str(exc))
            raise TypeError("eval_metrics must be an sklearn.eval_metrics class, or a function with the signature "
                            "(y_true, y_pred, **kwargs)'")
        except ValueError as exc:
            logger.exception(exc)
            raise ValueError("arguments not in the correct format for metric")

    def evaluate_generator_model(self, dataloader, eval_metrics=None, **kwargs):
        """
        evaluates the model based on the provided dataloader
        :param dataloader: a pytorch dataloader with test dataset and labels
        :type dataloader: torch.utils.data.DataLoader
        :param eval_metrics: a function for how to evaluate the batched examples and labels
        must have the signature (x_batch, y_batch) => float
        :type eval_metrics: 'function'
        :return: dictionary of metrics
        :rtype: 'dict'
        """
        metric_score = 0
        for x_batch, y_batch in dataloader:
            y_pred = self.model.evaluation_step(x_batch)

            if eval_metrics is None:
                y_pred = np.exp(y_pred)
                y_pred = np.argmax(y_pred, axis=1)
                equals = (y_pred == y_batch)
                metric_score += torch.mean(equals.type(torch.FloatTensor))

            else:
                try:
                    metric_score += eval_metrics(x_batch, y_batch)

                except TypeError as exc:
                    logger.exception(str(exc))
                    raise TypeError("eval_metrics must be a function "
                                    "with the signature (x_batch, y_batch, **kwargs)'")

        metric_dict = {'metric_score': metric_score / len(dataloader)}

        return metric_dict

    def save_model(self, filename=None, optimizer_filename=None, history_filename=None, onnx_export=True):
        """
        Save a model to file in the format specific to the backend framework.

        :param filename: Name of the file that contains the model to be loaded.
        :type filename: `str`
        :param optimizer_filename: Name of the file that contains the optimizer to be loaded.
        :type optimizer_filename: `str`
        :param history_filename: Name of the file that contains the model history to be loaded.
        :type history_filename: `str`
        :return: None
        """
        if filename is None:
            file = self.model_name if self.model_name else self.model_type
            if onnx_export:
                filename = '{}_{}.onnx'.format(file, time.time())
            else:
                filename = '{}_{}.pt'.format(file, time.time())

        if onnx_export:
            x = torch.randn(1,1,28,28)
            torch.onnx.export(self.model.module_, x, super().get_model_absolute_path(filename), export_params=True, keep_initializers_as_inputs=True)
        else:
            f_params = super().get_model_absolute_path(filename)
            f_optimizer = None
            f_history = None
            if optimizer_filename is not None:
                f_optimizer = super().get_model_absolute_path(optimizer_filename)
            if history_filename is not None:
                f_history = super().get_model_absolute_path(history_filename)
            self.model.save_params(
                f_params=f_params, f_optimizer=f_optimizer, f_history=f_history)

        return filename

    def load_model(self, pytorch_module, model_filename, optimizer_filename=None, history_filename=None, optimizer=None,
                   module_init_params=None):
        """
        Loads a model from disk given the specified file_name

        :param pytorch_module: uninstantiated pytorch model class
        :type pytorch_module: torch.nn.Module class reference
        :param model_filename: Name of the file that contains the model to be loaded.
        :type model_filename: `str`
        :param optimizer_filename: Name of the file that contains the optimizer to be loaded.
        :type optimizer_filename: `str`
        :param history_filename: Name of the file that contains the model history to be loaded.
        :type history_filename: `str`
        :param optimizer: the optimizer that should be loaded
        :type optimizer: pytorch optimizer class
        :param module_init_params: A dictionary with the values for the model's init arguments. \
        The key for the init parameters must be prefixed with module__ \
        (ex. if the parameter name is hidden_size, then key must be module__hidden_size)
        :type module_init_params: 'dict'
        :return: None
        """
        f_params = model_filename
        f_optimizer = None
        f_history = None
        if optimizer is not None and optimizer_filename is not None:
            self.optimizer = optimizer
            f_optimizer = optimizer_filename
        if history_filename is not None:
            f_history = history_filename
        model = self.initialize_model(pytorch_module, self.optimizer, self.criterion,
                                      module_init_params=module_init_params)
        model.load_params(f_params=f_params,
                          f_optimizer=f_optimizer, f_history=f_history)

        if self.use_gpu_for_training and torch.cuda.device_count() > 0:
            if self.num_gpus > torch.cuda.device_count():
                logger.error('Selected number of gpus to use for training exceeds number of available gpus, ' +
                             str(torch.cuda.device_count()) +
                             'Set number of gpus to maximum available on device.')
                self.num_gpus = torch.cuda.device_count()
            device_ids = list(range(self.num_gpus))
            model.module_ = torch.nn.DataParallel(model.module_, device_ids=device_ids)
            model.set_params(device='cuda')
            model.module_.to('cuda')

        self.model = model

    def get_weights(self, to_numpy=False):
        """
        Returns the weights of the model

        :param to_numpy; Determines whether the weights should be returned as numpy array, or tensor
        :type to_numpy: `boolean`
        :return: list of model weights
        """
        if self.use_gpu_for_training and torch.cuda.device_count() > 0:
            module = copy.deepcopy(self.model.module_).cpu()

        else:
            module = self.model.module_
        if to_numpy:
            return self.parameters_to_numpy(module.parameters())
        return list(module.parameters())

    def parameters_to_numpy(self, params):
        """
        Transforms parameter tensors to numpy arrays

        :param params: The parameter tensor to be transformed
        :return: numpy array of parameters 
        """
        np_params = []
        for layer in params:
            np_params.append(layer.detach().numpy())
        return np_params

    def load_model_from_spec(self, model_spec):
        """
        Loads model from provided model_spec, where model_spec is a `dict`
        that contains one item: model_spec['model_definition'], which has a
        pointer to the file where an nn.sequence container is saved

        :return: model
        :rtype: `nn.sequence`
        """

        model_file = model_spec['model_definition']
        model_absolute_path = config.get_absolute_path(model_file)
        model = torch.load(model_absolute_path)
        return model

    def get_gradient(self, train_data):
        """
        Returns the gradients for each layer in the model

        :return: gradients
        :rtype: `list`. numpy array list of model's gradients
        """

        x = train_data[0]
        y = train_data[1]
        gradients = []
        y_infer = self.model.module_(torch.from_numpy(x))
        loss = self.model.criterion_(y_infer, torch.from_numpy(y))
        loss.backward()

        for layer in self.model.module_.parameters():
            gradients.append(layer.grad.numpy())

        return gradients

    @staticmethod
    def valid_acc(net, x, y):
        """
        Callback scoring function that returns the validation accuracy during training.
        :param net: The model that will be used
        :param x: the validation data set
        :param y: the training targets
        :return: the accuracy of the validation training pass
        """
        y_pred = net.predict(x)
        y_pred = np.exp(y_pred)
        y_pred = np.argmax(y_pred, axis=1)
        return metrics.accuracy_score(y, y_pred)

    @staticmethod
    def loss(y, y_pred, net):
        loss = net.get_loss(torch.from_numpy(y_pred), torch.from_numpy(y)).numpy()

        return loss

    @staticmethod
    def __expand_linear_layer__(net, layer_name, new_dimensions, layer_idx):
        """
        Expands a linear layer of a pytorch module
        :param net: an instantiated pytorch module
        :type net: `torch.nn.Module`
        :param layer_name: layer variable name which needs to be expanded
        :type layer_name: `str`
        :param new_dimensions: new dimensions
        :type new_dimensions: `list`
        :param layer_idx: layer idx which needs to be expanded. this corresponds with
        the index of `new_dimensions`
        :type layer_idx: `int`
        :return: None
        """

        original_lin_layer = getattr(net, layer_name)

        if not isinstance(original_lin_layer, torch.nn.Linear):
            raise FLException('Received a non-linear layer to expand '
                              'whereas the method expects a linear layer')

        new_ip_dim = original_lin_layer.in_features if layer_idx == 0 else new_dimensions[
            layer_idx-1]
        new_op_dim = new_dimensions[layer_idx]
        bias = original_lin_layer.bias is not None

        new_lin_layer = type(original_lin_layer)(
            in_features=new_ip_dim, out_features=new_op_dim, bias=bias
        )

        setattr(net, layer_name, new_lin_layer)

    def expand_model_by_layer_name(self, new_dimension, layer_name="dense"):
        """
        Expands the current PyTorch models layers with the provided dimensions
        :param new_dimension: new dimensions of the particular `layer_name`
        :type new_dimension: `list`
        :param layer_name: layer name which needs to be expanded
        :type layer_name: `str`
        :return: None
        """

        if new_dimension is None:
            raise FLException('No information is provided for '
                              'the new expanded model. '
                              'Please provide the new dimension of '
                              'the resulting expanded model')

        layer_maps = {
            'dense': {
                'class': torch.nn.Linear,
                'expansion_fn': self.__expand_linear_layer__
            }
        }

        layer_cls = layer_maps[layer_name]['class']
        layer_exp_fn = layer_maps[layer_name]['expansion_fn']
        net = self.model.module_    # get instantiated module
        idx = 0

        for layer_varname, layer in net.named_modules():
            # skip layer if not instance of layer desired to be expanded
            if not isinstance(layer, layer_cls):
                continue

            layer_exp_fn(net, layer_varname, new_dimension, idx)
            idx += 1

        self.model.initialize_optimizer()

    def is_fitted(self):
        """
        Return a boolean value indicating if the model is fitted or not.
        In particular, it calls `skorch.utils.check_is_fitted` to checks
        whether the net is initialized.
        If it has, return True; otherwise return false.

        :return: res
        :rtype: `bool`
        """
        try:
            self.model.check_is_fitted()
        except NotInitializedError:
            return False
        return True

    def get_loss(self, dataset):
        """
        Return the resulting loss computed based on the provided dataset.

        :param train_data: Training data, a tuple \
        given in the form (x_train, y_train). otherwise, input compatible with skorch.dataset.Dataset
        :type train_data: `(np.ndarray, ”np.ndarray)`
        :return: The resulting loss.
        :rtype: `float`
        """
        
        metrics_list = [PytorchFLModel.loss]
        orig_loss_value = self.evaluate(dataset, eval_metrics=metrics_list)[
                    'loss']
        
        return orig_loss_value
