import numpy as np
from matplotlib import cm
import matplotlib.pyplot as plt
from smt.problems import Rosenbrock
from smt.sampling_methods import LHS

########### Initialization of the problem, construction of the training and validation points

ndim = 2
ndoe = 20  # int(10*ndim)

# Define the function
fun = Rosenbrock(ndim=ndim)

# Construction of the DOE
# in order to have the always same LHS points, random_state=1
sampling = LHS(xlimits=fun.xlimits, criterion='ese', random_state=1)
xt = sampling(ndoe)
# Compute the outputs
yt = fun(xt)

# Construction of the validation points
ntest = 200 #500
sampling = LHS(xlimits=fun.xlimits, criterion='ese', random_state=1)
xtest = sampling(ntest)
ytest = fun(xtest)

# To plot the Rosenbrock function
x = np.linspace(-2, 2, 50)
res = []
for x0 in x:
    for x1 in x:
        res.append(fun(np.array([[x0, x1]])))
res = np.array(res)
res = res.reshape((50, 50)).T
X, Y = np.meshgrid(x, x)
fig = plt.figure(figsize=(6, 4))
ax = fig.add_subplot(projection="3d")
surf = ax.plot_surface(
    X, Y, res, cmap=cm.viridis, linewidth=0, antialiased=False, alpha=0.5
)

# Bad statements
ax.scatter(xt[:, 0], xt[:, 1], yt, zdir="z", marker="x", c="b", s=200, label="Training point")
ax.scatter(xtest[:, 0], xtest[:, 1], ytest, zdir="z", marker=".", c="k", s=200, label="Validation point")

# Updated statements
# ax.scatter(xt[:, 0], xt[:, 1], yt[:, 0], zdir='z', marker='x', c='b', s=200, label='Training point')
# ax.scatter(xtest[:, 0], xtest[:, 1], ytest[:, 0], zdir='z', marker='.', c='k', s=200, label='Validation point')

plt.title("Rosenbrock function")
plt.xlabel("x1")
plt.ylabel("x2")
plt.legend()
plt.show()