import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def psi(x):
    return np.exp(x) - 1

def phi(x, N):
    intervals = np.linspace(0, 1, 5*N+1)
    values = []
    for i in range(N):
        x_min = intervals[5*i]
        x_max = intervals[5*i+4]
        constant_value = np.floor(50 * psi(x_min)) / 50
        values.extend([constant_value] * 5)
    values.append(values[-1])  # Add last value to make it the same length as intervals
    return np.interp(x, intervals, values)
N = 10

# Create figure with two subplots
fig = plt.figure(figsize=(15, 6))
ax1 = fig.add_subplot(121)

# Subplot 1: psi(x) and phi(x)
x = np.linspace(0, 1, 1000)
ax1.plot(x, psi(x), label='ψ(x)')
ax1.plot(x, phi(x, N), label='φ(x)', linestyle='--')
# ax1.set_xlabel('x')
# ax1.set_ylabel('y')
ax1.legend()
# ax1.set_title('ψ(x) and φ(x)')

# Subplot 2: 3D plot of phi(x) + t*phi(y)
x = y = np.linspace(0, 1, 100)
X, Y = np.meshgrid(x, y)
t = (np.sqrt(5) - 1) / 2
Z = phi(X, N) + t * phi(Y, N)

ax2 = fig.add_subplot(122, projection='3d')
ax2.plot_surface(X, Y, Z, cmap='viridis')
ax2.set_xlabel('x')
ax2.set_ylabel('y')
# ax2.set_zlabel('φ(x) + t·φ(y)')
ax2.set_title('φ(x) + t·φ(y) on [0, 1]²')
# remove axes
# ax2.set_axis_off()

plt.tight_layout()
plt.savefig('Kolmogorov-Arnold grid.svg')
plt.show()