from re import X
import numpy as np
import matplotlib.pyplot as plt

def f(x):
  return np.sin(2 * np.pi * x) * np.exp(-x)

X_train = np.array([0.0, 0.4, 0.7, 1.0, 1.5, 2.0])
y_train = f(X_train)

def rbf(r, epsilon=2.0):
  # Corrected to return an array of the same shape as r
  return np.exp(-(epsilon * r)**2)

def rbf_model(X, y, epsilon=2.0):
  R = np.abs(X[:, None] - X[None, :])
  Phi = rbf(R, epsilon)
  weights = np.linalg.solve(Phi, y)
  return weights

def rbf_predict(X_new, X_train, weights, epsilon=2.0):
  R = np.abs(X_new[:, None] - X_train[None, :])
  Phi = rbf(R, epsilon)
  return Phi @ weights

weights = rbf_model(X_train, y_train)
X_new = np.linspace(0, 2, 100)
y_pred = rbf_predict(X_new, X_train, weights)
# Corrected X_test to X_new
y_true = f(X_new)

plt.plot(X_new, y_true, label='True Function')
plt.plot(X_new, y_pred, '--', label='RBF surrogate')
plt.scatter(X_train, y_train, color='red', label='Training Points')
plt.title('RBF Surrogate Model')
plt.xlabel('X')
plt.ylabel('f(x)')
plt.legend()
plt.grid(True)
plt.show()