import numpy as np
from sklearn.gaussian_process import GaussianProcessRegressor


def f_demo_1d(*args, **kwargs):
    coord = args[0]
    X = coord[0]
    return np.exp(-1/(1+X))*(np.sin(2*X) + np.sin(5*np.pi*np.log(1 + X)**2))


def f_demo_2d(*args, **kwargs):
    coord = args[0]
    XX = coord[0]
    YY = coord[1]
    return np.sin(XX) + np.cos(YY)

def f_demo_3d(*args, **kwargs):
    coord = args[0]
    XX = coord[0]
    YY = coord[1]
    ZZ = coord[2]
    return np.sin(XX) + np.cos(YY) + np.cos(ZZ)

def f_synthetic_test_chowdhury_gopalan(x, kernel, n_samples):
    x = x[:, np.newaxis]
    gp = GaussianProcessRegressor(kernel=kernel)

    y_mean = gp.predict(x)
    y_samples = gp.sample_y(x, n_samples)

    range_y = np.array([y_samples[:, n].max() - y_samples[:, n].min() for n in range(n_samples)])
    R = np.sqrt(0.01*range_y)
    lambda_param = R**2

    kernel_matrix = kernel(x, x)
    B = np.array([np.sqrt(np.dot(np.dot(y_samples[:, n], kernel_matrix), y_samples[:, n])) for n in range(n_samples)])
    return np.rollaxis(y_samples, 1), R, lambda_param, B

def f_rosenbrock(*args, **kwargs):
    coord = args[0]
    XX = coord[0]
    YY = coord[1]
    return (1 - XX)**2 + 100*(YY - XX**2)**2

def f_rosenbrock_min(*args, **kwargs):
    return - f_rosenbrock(*args, **kwargs)

def f_hartman3(*args, **kwargs):
    coord = args[0]
    XX = coord[0]
    YY = coord[1]
    ZZ = coord[2]
    A = np.array([[3., 10, 30],
                  [0.1, 10, 35],
                  [3., 10, 30],
                  [0.1, 10, 35]])
    alpha = np.array([1.0, 1.2, 3.0, 3.2])
    P = 1e-4*np.array([[3689, 1170, 2673],
                       [4699, 4387, 7470],
                       [1091, 8732, 5547],
                       [381, 5743, 8828]])
    return -sum(alpha[i]*np.exp(-sum(A[i,j]*(coord[j] - P[i,j])**2 for j in range(3))) for i in range(4))

def f_hartman3_min(*args, **kwargs):
    return - f_hartman3(*args, **kwargs)

def make_params(x, f, kernel):
    y = f(x)
    range_y = np.linalg.norm(np.array([y.max() - y.min()]))
    R = np.sqrt(0.01*range_y)
    lambda_param = R**2
    return R, lambda_param
