import numpy as np
from profit.sur import Surrogate
from profit.sur.gp import GaussianProcess
from profit.defaults import fit_gaussian_process as defaults, fit as base_defaults


@Surrogate.register('GPy')
class GPySurrogate(GaussianProcess):
    """Surrogate for https://github.com/SheffieldML/GPy.

    Attributes:
        model (GPy.models): Model object of GPy.
    """

    def __init__(self):
        import GPy
        self.GPy = GPy
        super().__init__()
        self.model = None

    def train(self, X, y, kernel=defaults['kernel'], hyperparameters=defaults['hyperparameters'],
              fixed_sigma_n=base_defaults['fixed_sigma_n']):
        self.pre_train(X, y, kernel, hyperparameters, fixed_sigma_n)
        self.model = self.GPy.models.GPRegression(self.Xtrain, self.ytrain, self.kernel,
                                                  noise_var=self.hyperparameters['sigma_n'] ** 2)
        self.optimize()
        self.post_train()

    def add_training_data(self, X, y):
        """Adds training points to the existing dataset.

        This is important for Active Learning. The data is added but the hyperparameters are not optimized yet.

        Parameters:
            X (ndarray): Input points to add.
            y (ndarray): Observed output to add.
        """
        self.Xtrain, self.ytrain = np.concatenate([self.Xtrain, X], axis=0), np.concatenate([self.ytrain, y], axis=0)
        self.encode_training_data()
        self.model.set_XY(self.Xtrain, self.ytrain)
        self.decode_training_data()

    def set_ytrain(self, y):
        """Set the observed training outputs. This is important for active learning.

        Parameters:
        y (np.array): Full training output data.
        """
        self.ytrain = np.atleast_2d(y.copy())
        self.encode_training_data()
        self.model.set_XY(self.Xtrain, self.ytrain)
        self.decode_training_data()

    def predict(self, Xpred, add_data_variance=True):
        Xpred = self.pre_predict(Xpred)
        ymean, yvar = self.model.predict(Xpred, include_likelihood=add_data_variance)
        ymean, yvar = self.decode_predict_data(ymean, yvar)
        return ymean, yvar

    def save_model(self, path):
        """Save the model as dict to a .hdf5 file.

        Parameters:
            path (str): Path including the file name, where the model should be saved.
        """

        from profit.util.file_handler import FileHandler

        sur_dict = {attr: getattr(self, attr) for attr in ('Xtrain', 'ytrain')}
        sur_dict['model'] = self.model.to_dict()
        sur_dict['input_encoders'] = str([enc.repr for enc in self.input_encoders])
        sur_dict['output_encoders'] = str([enc.repr for enc in self.output_encoders])
        FileHandler.save(path, sur_dict)

    @classmethod
    def load_model(cls, path):
        """Loads a saved model from a .hdf5 file and updates its attributes. In case of a multi-output model, the .pkl
        file is loaded, since .hdf5 is not supported yet.

        Parameters:
            path (str): Path including the file name, from where the model should be loaded.

        Returns:
            GPy.models: Instantiated surrogate model.
        """

        from profit.util.file_handler import FileHandler
        from profit.sur.encoders import Encoder
        from GPy import models
        from numpy import array  # needed for eval of arrays

        self = cls()
        sur_dict = FileHandler.load(path, as_type='dict')
        self.model = models.GPRegression.from_dict(sur_dict['model'])
        self.Xtrain, self.ytrain = sur_dict['Xtrain'], sur_dict['ytrain']

        for enc in eval(sur_dict['input_encoders']):
            self.add_input_encoder(Encoder[enc['class']](enc['columns'], enc['parameters']))
        for enc in eval(sur_dict['output_encoders']):
            self.add_output_encoder(Encoder[enc['class']](enc['columns'], enc['parameters']))

        # Initialize the encoder by encoding and decoding the training data once.
        self.encode_training_data()
        self.decode_training_data()

        self.kernel = self.model.kern
        self.ndim = self.Xtrain.shape[-1]
        self._set_hyperparameters_from_model()
        self.trained = True
        self.print_hyperparameters("Loaded")
        return self

    def select_kernel(self, kernel):
        """Get the GPy.kern kernel by matching the given string kernel identifier.

        Parameters:
            kernel (str): Kernel string such as 'RBF' or depending on the surrogate also product and sum kernels
                such as 'RBF+Matern52'.

        Returns:
            GPy.kern: GPy kernel object. Currently, for sum and product kernels, the initial hyperparameters are the
            same for all kernels.
        """

        try:
            if not any(operator in kernel for operator in ('+', '*')):
                # Single kernel
                return getattr(self.GPy.kern, kernel)(self.ndim,
                                                      lengthscale=self.hyperparameters.get('length_scale', [1]),
                                                      variance=self.hyperparameters.get('sigma_f', 1)**2,
                                                      ARD=len(self.hyperparameters.get('length_scale', [1])) > 1)
            else:
                from re import split
                full_str = split('([+*])', kernel)
                kern = []
                for key in full_str:
                    kern += [key if key in ('+', '*') else
                               'self.GPy.kern.{}({}, lengthscale={}, variance={})'
                                   .format(key, self.ndim,
                                           self.hyperparameters.get('length_scale', [1]),
                                           self.hyperparameters.get('sigma_f', 1)**2)]
                return eval(''.join(kern))
        except AttributeError:
            raise RuntimeError("Kernel {} is not implemented.".format(kernel))

    def optimize(self, return_hess_inv=False, **opt_kwargs):
        """For hyperparameter optimization the GPy base optimization is used.

        Currently, the inverse Hessian can not be retrieved, which limits the active learning effectivity.

        Parameters:
            return_hess_inv (bool): Is not considered currently.
            opt_kwargs: Keyword arguments used directly in the GPy base optimization.
        """

        self.model.optimize(**opt_kwargs)
        self._set_hyperparameters_from_model()
        self.print_hyperparameters("Optimized")

    def _set_hyperparameters_from_model(self):
        r"""Helper function to set the hyperparameter dict from the model.

        It depends on whether it is a single kernel or a combined one.
        """
        if hasattr(self.model.kern, 'lengthscale'):
            self.hyperparameters['length_scale'] = self.model.kern.lengthscale.values
            self.hyperparameters['sigma_f'] = np.sqrt(self.model.kern.variance)
            self.hyperparameters['sigma_n'] = np.sqrt(self.model.likelihood.variance)
        elif hasattr(self.model.kern, 'parts'):
            for part in self.model.kern.parts:
                for key, value in zip(part.parameter_names(), part.param_array):
                    value = np.atleast_1d(value)
                    if key == 'lengthscale':
                        self.hyperparameters['length_scale'] = value
                    elif key == 'variance':
                        self.hyperparameters['sigma_f'] = np.sqrt(value)
                    else:
                        self.hyperparameters[key] = value
            noise_var = self.model.likelihood.gaussian_variance(
                self.model.Y_metadata).reshape(-1, 1, order='F')
            self.hyperparameters['sigma_n'] = np.sqrt(np.max(noise_var, axis=0))
        self.decode_hyperparameters()

    def special_hyperparameter_decoding(self, key, value):
        has_ard = any(p.ARD for p in self.model.kern.parts) if hasattr(self.model.kern, 'parts') else self.model.kern.ARD
        if key == 'length_scale' and self.ndim > 1 and not has_ard:
            return np.atleast_1d(np.linalg.norm(value))
        elif key in ('sigma_f', 'sigma_n') and len(value) > 1:
            return np.atleast_1d(np.max(value))
        return value


@Surrogate.register("CoregionalizedGPy")
class CoregionalizedGPySurrogate(GPySurrogate):

    def pre_train(self, X, y, kernel=defaults['kernel'], hyperparameters=defaults['hyperparameters'],
                  fixed_sigma_n=base_defaults['fixed_sigma_n']):
        super().pre_train(X, y, kernel, hyperparameters, fixed_sigma_n)
        self.output_ndim = self.ytrain.shape[-1]

    def train(self, X, y, kernel=defaults['kernel'], hyperparameters=defaults['hyperparameters'],
              fixed_sigma_n=base_defaults['fixed_sigma_n']):
        self.pre_train(X, y, kernel, hyperparameters, fixed_sigma_n)
        icm = self.GPy.util.multioutput.ICM(input_dim=self.ndim, num_outputs=self.output_ndim, kernel=self.kernel)
        _X = self.output_ndim * [self.Xtrain]
        _y = [self.ytrain[:, d].reshape(-1, 1) for d in range(self.output_ndim)]
        self.model = self.GPy.models.GPCoregionalizedRegression(_X, _y, kernel=icm)
        self.optimize()
        self.post_train()

    def add_training_data(self, X, y):
        self.Xtrain, self.ytrain = np.concatenate([self.Xtrain, X], axis=0), np.concatenate([self.ytrain, y], axis=0)
        self.encode_training_data()
        new_Xtrain = np.empty((self.Xtrain.shape[0] * self.output_ndim, self.Xtrain.shape[-1] + 1))
        for d in range(self.output_ndim):
            start_idx = self.Xtrain.shape[0] * d
            end_idx = start_idx + self.Xtrain.shape[0]
            new_Xtrain[start_idx:end_idx] = np.hstack([self.Xtrain, np.ones((self.Xtrain.shape[0], 1)) * d])
        new_noise_dict = {'output_index': new_Xtrain[:, -1:].astype(int)}
        self.model.Y_metadata = new_noise_dict
        self.model.set_XY(new_Xtrain, self.ytrain.reshape(-1, 1, order='F'))
        self.decode_training_data()

    def set_ytrain(self, y):
        self.ytrain = np.atleast_2d(y.copy())
        self.encode_training_data()
        self.model.set_Y(self.ytrain.reshape(-1, 1, order='F'))
        self.decode_training_data()

    def predict(self, Xpred, add_data_variance=True):
        Xpred = super().pre_predict(Xpred)
        ymean = np.empty((Xpred.shape[0], self.output_ndim))
        yvar = ymean.copy()
        for d in range(self.output_ndim):
            newX = np.hstack([Xpred, np.ones((Xpred.shape[0], 1)) * d])
            noise_dict = {'output_index': newX[:, -1:].astype(int)}
            ym, yv = self.model.predict(newX, Y_metadata=noise_dict)
            ymean[:, d], yvar[:, d] = ym.flatten(), yv.flatten()
        ymean, yvar = self.decode_predict_data(ymean, yvar)
        return ymean, yvar

    def save_model(self, path):
        # GPy does not support to_dict method for Coregionalization kernel yet.
        from profit.util.file_handler import FileHandler
        from os.path import splitext
        filepath, ending = splitext(path)
        if ending != '.pkl':
            print(f"Saving to '{ending}' not implemented yet for {self.__class__.__name__} surrogate."
                  " Saving to '.pkl' file instead.")
        save_list = [self.model, self.Xtrain, self.ytrain,
                     str([enc.repr for enc in self.input_encoders]),
                     str([enc.repr for enc in self.output_encoders])]
        FileHandler.save(filepath + '.pkl', save_list)

    @classmethod
    def load_model(cls, path):
        from profit.util.file_handler import FileHandler
        from profit.sur.encoders import Encoder
        from numpy import array  # needed for eval of arrays

        self = cls()
        try:
            self.model, self.Xtrain, self.ytrain, input_encoder_str, output_encoder_str = FileHandler.load(path, as_type='raw')
        except (OSError, FileNotFoundError):
            from os.path import splitext
            print("File not found. Try changing file ending to '.pkl'!")
            path = splitext(path)[0] + '.pkl'
            # Load multi-output model from pickle file
            self.model, self.Xtrain, self.ytrain, input_encoder_str, output_encoder_str = FileHandler.load(path, as_type='raw')

        for enc in eval(input_encoder_str):
            self.add_input_encoder(Encoder[enc['class']](enc['columns'], enc['parameters']))
        for enc in eval(output_encoder_str):
            self.add_output_encoder(Encoder[enc['class']](enc['columns'], enc['parameters']))
        self.output_ndim = int(max(self.model.X[:, -1])) + 1

        # Initialize the encoder by encoding and decoding the training data once.
        self.encode_training_data()
        self.decode_training_data()

        self.kernel = self.model.kern
        self.ndim = self.Xtrain.shape[-1]
        self._set_hyperparameters_from_model()
        self.trained = True
        self.print_hyperparameters("Loaded")
        return self

    def special_hyperparameter_decoding(self, key, value):
        if key == 'length_scale' and self.ndim > 1 and not self.model.kern.parts[0].ARD:
            return np.atleast_1d(np.linalg.norm(value))
        elif key in ('sigma_f', 'sigma_n') and len(value) > 1:
            return np.atleast_1d(np.max(value))
        elif key in ('W', 'kappa'):
                new_value = value
                for enc in self.output_encoders:
                    if enc.label == 'Normalization':
                        new_value = enc.decode_hyperparameters(value)
                return np.atleast_1d(np.min(new_value))
        return value
