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

# Function to plot a surface z = sin(x^2 + y^2) / (x^2 + y^2)
def plot_surface(ax):
    x = np.linspace(-5, 5, 100)
    y = np.linspace(-5, 5, 100)
    X, Y = np.meshgrid(x, y)
    Z = (np.sin(X**2 + Y**2)) / (X**2 + Y**2)
    
    ax.plot_surface(X, Y, Z, cmap='viridis', edgecolor='none')
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.set_title(r'$Z = \frac{\sin(X^2 + Y^2)}{X^2 + Y^2}$')


# Function to plot the Klein bottle using parametric equations
def plot_klein_bottle(ax):
    t = np.linspace(0, 2 * np.pi, 100)
    v = np.linspace(0, 2 * np.pi, 100) # add parameter for v to define the Klein bottle surface
    T, V = np.meshgrid(t, v) # create a grid of t and v

    x = np.sin(V) * (np.cos(T) + 2)
    y = np.cos(V) * (np.cos(T) + 2)
    z = np.sin(T) * np.cos(V)


    # Plot the surface with color and smoothing parameters
    ax.plot_surface(x, y, z, color='b', rstride=4, cstride=4) # Adjust rstride and cstride for smoother or more detailed surface
    ax.set_xlim([-2, 2])
    ax.set_ylim([-2, 2])
    ax.set_zlim([-1, 1])
    ax.set_title('Klein Bottle')


# Main function to call the plotting functions
def main():
    fig = plt.figure(figsize=(12, 6))  # Create a figure with desired size

    ax1 = fig.add_subplot(121, projection='3d') # First subplot (3D)
    plot_surface(ax1) # Pass ax1 to plot the surface

    ax2 = fig.add_subplot(122, projection='3d') # Second subplot (3D)
    plot_klein_bottle(ax2) # Pass ax2 to plot the Klein bottle

    plt.tight_layout()  # Adjust spacing between subplots

    plt.show() 

if __name__ == "__main__":
    main()