import numpy as np
from sklearn.preprocessing import MinMaxScaler
import random
from RatBrainModel import RatEst
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import RBF, ConstantKernel as C
from sklearn.gaussian_process.kernels import DotProduct, WhiteKernel
from sklearn.gaussian_process.kernels import (RBF, Matern, RationalQuadratic,
                                              ExpSineSquared, DotProduct,
                                              ConstantKernel)
from matplotlib import pyplot as plt

import pandas as pd

from scipy.stats import norm

from scipy.optimize import minimize

#from bayesian_optimization_util import plot_approximation, plot_acquisition


def propose_location(acquisition, X_sample, Y_sample, gpr, bounds, n_restarts=25):
    '''
    Proposes the next sampling point by optimizing the acquisition function.

    Args:
        acquisition: Acquisition function.
        X_sample: Sample locations (n x d).
        Y_sample: Sample values (n x 1).
        gpr: A GaussianProcessRegressor fitted to samples.

    Returns:
        Location of the acquisition function maximum.
    '''
    dim = X_sample.shape[1]
    min_val = 1
    min_x = None

    def min_obj(X):
        # Minimization objective is the negative acquisition function
        return -acquisition(X.reshape(-1, dim), X_sample, Y_sample, gpr)

    # Find the best optimum by starting from n_restart different random points.
    for x0 in np.random.uniform(bounds[:, 0], bounds[:, 1], size=(n_restarts, dim)):
        res = minimize(min_obj, x0=x0, bounds=bounds, method='L-BFGS-B')
        if res.fun < min_val:
            min_val = res.fun[0]
            min_x = res.x

    return min_x.reshape(-1, 1)


def expected_improvement(X, X_sample, Y_sample, gpr, xi=0.6):
    '''
    Computes the EI at points X based on existing samples X_sample
    and Y_sample using a Gaussian process surrogate model.

    Args:
        X: Points at which EI shall be computed (m x d).
        X_sample: Sample locations (n x d).
        Y_sample: Sample values (n x 1).
        gpr: A GaussianProcessRegressor fitted to samples.
        xi: Exploitation-exploration trade-off parameter.

    Returns:
        Expected improvements at points X.
    '''
    mu, sigma = gpr.predict(X, return_std=True)
    #mu_sample = gpr.predict(X_sample)

    sigma = sigma.reshape(-1, 1)

    # Needed for noise-based model,
    # otherwise use np.max(Y_sample).
    # See also section 2.4 in [...]
    mu_sample_opt = np.max(Y_sample)

    with np.errstate(divide='warn'):
        imp = mu - mu_sample_opt - xi
        Z = imp / sigma
        ei = imp * norm.cdf(Z) + sigma * norm.pdf(Z)
        ei[sigma == 0.0] = 0.0

    return ei


# Update Gaussian process with existing samples
def acquire_bayesian_ei(gpEst,X_sample,Y_sample):
    bounds = np.array([[0, 4.0]])
    print(X_sample.shape)
    print(Y_sample.shape)
    # Obtain next sampling point from the acquisition function (expected_improvement)
    X_next = propose_location(expected_improvement, X_sample, Y_sample, gpEst, bounds)

    # Obtain next noisy sample from the objective function
    return X_next

def acquire_bayesian_uncertainty(gpEst):
    x_num=50
    y_num=20
    x = np.atleast_2d(np.linspace(0, 4, x_num)).T

    """y_pred= gpEst.sample_y(x,n_samples=y_num,random_state=None)
    y_pred = np.array(y_pred)
    y_pred=y_pred.reshape(x_num,y_num)
    y_std = np.std(y_pred, axis=1)
    """

    y_pred,sigma = gpEst.predict(x, return_std=True)
    sigma = np.array(sigma)
    i=np.argmax(sigma)
    return x[i]



"""
gp.fit(X, y)

x = np.atleast_2d(np.linspace(0,4 , 10)).T

y_pred, sigma = gp.predict(x, return_std=True)

###############Estimation Model
num_samples=5
X_samples=np.random.uniform(low=0.0, high=4.0, size=(num_samples,1))
Y_samples=gp.sample_y(X_samples)
print(x)
print(Y_samples)
"""