import numpy as np
from numpy import diag, log, ones
from numpy import concatenate as cat
from scipy.linalg import cholesky, cho_solve
from scipy.optimize import minimize
from .optim import Optimizer, gridSearch
import warnings

class GaussianProcess():
    def __init__(self, kernel):
        """This class implements a Gaussian Process (GP).
        
        Parameters
        ----------

        kernel : an object of a kernel class
        """

        self.kernel = kernel
        self.x_train = None
        self.y_train = None
        self.y_var_train = None
        self.var_n = 1e-10
        self.L = None
        self.alpha = None
        self.is_bad_theta = False
    
    def addData(self, x_train, y_train, y_var_train=None, replace=False):
        """This method adds training data to the Gaussian Process with the option 
        to specify the sampling noise and to replace existing data.

        Parameters
        ----------

        x_train : (n_samples, n_features) np.array

        y_train : (n_samples, 1) np.array

        y_var_train : (n_samples, 1) np.array

        replace : bool
        """
        # delete training data, if replace is True
        if replace:
            self.x_train, self.y_train, self.y_var_train = None, None, None
        
        # add x_train and y_train to properties
        self.x_train = x_train if self.x_train is None else cat((self.x_train, x_train), axis=0)
        self.y_train = y_train if self.y_train is None else cat((self.y_train, y_train), axis=0)

        # ensure var_n is a 1D array of same length as y_train and set to zeros, if not specified
        if not y_var_train is None:
            y_var_train = np.array(y_var_train)
            y_var_train = y_var_train.flatten()*ones(self.y_train.size)         
            self.y_var_train = y_var_train if self.y_var_train is None else cat(self.y_var_train, y_var_train, axis=0)

    def predict(self, x_pred, predict_var=False):
        """Inference method of the Gaussian Process.
        
        Parameters
        ----------

        x_pred : (n_samples, n_features) np.array

        predict_var : bool

        Returns
        -------

        y_mean : (n_samples, 1) np.array
                 predictive mean of the Gaussian Process
        
        y_var : (n_samples, 1) np.array
                 predictive variance of the Gaussian Process
        """
        k = self.kernel(x_pred, self.x_train)
        y_mean = k@self.alpha
        
        if predict_var:
            v = cho_solve((self.L, True), k.T)            
            y_var = self.kernel.diag(x_pred) - np.einsum("ij,ji->i", k, v)
            y_var_negative = y_var < 0
            if np.any(y_var_negative):
                warnings.warn("Predicted variances smaller than 0. "
                                "Setting those variances to 0.")
                y_var[y_var_negative] = 0.0
            y_var = y_var.reshape(-1, 1)
            return y_mean, y_var
        return y_mean  

    def __call__(self, x_pred, predict_var=False):
        return self.predict(x_pred, predict_var)

    def updateWeights(self, theta=None):
        """This method computes K, L, and alpha based on 
        the training data and the hyperparameters theta.
        
        Parameters
        ----------

        theta : (n_params,) np.array
                hyperparameters of the kernel function
        """
        params = self.kernel.params if theta is None else theta.flatten()
        try:
            self.kernel.params = params
            self.K = self.kernel(self.x_train)
            np.fill_diagonal(self.K, self.K.diagonal() + self.var_n)
            if not self.y_var_train is None:
                np.fill_diagonal(self.K, self.K.diagonal() + self.y_var_train)
            self.L = cholesky(self.K, lower=True)
            self.alpha = cho_solve((self.L, True), self.y_train)
        except:
            self.is_bad_theta = True

    def fit(self, optim='grid_and_mincon', **kwargs):
        """This method optimizes the hyperparameters of the Gaussian Process.
        
        Parameters
        ----------

        optim : str
                optimization strategy for hyperparameter optimization 
                ['rand_starts', 'grid_starts', 'grid_and_mincon']

        N : int or list of ints
                Number of grid points per dimension in hyperparameter space
                or number of random starts in case of optim='rand_starts'
        
        logarithmic : boolean or list of booleans of length n_dim
                specify for each dimension whether points between bounds should 
                be distributed on log scale
        
        method : str
                gradient based optimization method for use in 
                scipy.optimize.minimize function (default is 'L-BFGS-B').
        
        seed : int
                seed of the random generator for random initial starts.
        """

        bounds = self.kernel.bounds
        lb, ub = bounds[:, 0], bounds[:, 1]
        
        theta = Optimizer.optimize(optimizer=optim, 
                           fun=self.negativeLogMarginalLikelihood,
                           lb=lb,
                           ub=ub,
                           **kwargs)
        # update weights
        self.updateWeights(theta)

    def negativeLogMarginalLikelihood(self, theta):
        """Loss function of the Gaussian Process for optimization of the hyperparameters.
        
        Parameters
        ----------

        theta : (n_params,) np.array
                hyperparameters of the kernel function

        Returns
        -------
        nlml : (1,) np.array

        """
        self.updateWeights(theta)
        if self.is_bad_theta:
            nlml = np.array([1e6])
            self.is_bad_theta = False
        else:
            nlml = 0.5*self.y_train.T@self.alpha + log(diag(self.L)).sum()
        return nlml.flatten()