from KDEpy import FFTKDE
func = lambda x : np.sin(x * 2 * np.pi) + (x + 1)**2

# Generate random data
num_data_points = 2**6
data_x = np.sort(np.random.rand(num_data_points))
data_y = func(data_x) + np.random.randn(num_data_points) / 5

# Plot the true function and the sampled values
x_smooth = np.linspace(0, 1, num=2**10)
plt.plot(x_smooth, func(x_smooth), label='True function')
plt.scatter(data_x, data_y, label='Function samples', s=10)

# Grid points in the x and y direction
grid_points_x, grid_points_y = 2**10, 2**4

# Stack the data for 2D input, compute the KDE
data = np.vstack((data_x, data_y)).T
kde = FFTKDE(bw=0.025).fit(data)
grid, points = kde.evaluate((grid_points_x, grid_points_y))

# Retrieve grid values, reshape output and plot boundaries
x, y = np.unique(grid[:, 0]), np.unique(grid[:, 1])
z = points.reshape(grid_points_x, grid_points_y)
plt.axvline(np.min(data_x), ls='--', c='k', label='Domain lower bound')
plt.axvline(np.max(data_x), ls='--', c='k', label='Domain upper bound')

# Compute y_pred = E[y | x] = sum_y p(y | x) * y
y_pred =  np.sum((z.T / np.sum(z, axis=1)).T  * y , axis=1)
plt.plot(x, y_pred, zorder=25, label='Kernel regression esimate')

plt.legend(); plt.tight_layout()