# surface3D.py
# https://matplotlib.org/stable/plot_types/3D/surface3d_simple.html#sphx-glr-plot-types-3d-surface3d-simple-py
'''
  The user selects between 7 functions to generate a 3D mesh
'''



import numpy as np
import matplotlib.pyplot as plt
 

# 7 functons for generating z coordinates
def fn1(x, y):
  r = np.sqrt(x*x + y*y)
  return -np.cos(r)

def fn2(x, y):
  r = np.sqrt(x*x + y*y)
  return np.exp(-r*r*0.1)

def fn3(x, y):
  r = np.sqrt(x*x + y*y)
  return np.sin(r)/r

def fn4(x, y):
  return np.sin(x)*np.cos(y)

def fn5(x, y):
  return np.cos(x*y*0.1)

def fn6(x, y):   # monkey saddle
  r = np.sqrt(x*x + y*y)
  return x*y*(x-y)*(x+y)/r

def fn7(x, y):
  return np.cos(0.1*x*np.exp(-y/5))


# ------------------------

n = int(input("n(1-7)? "))
if n < 1 or n > 7:
  n = 1

fn =  globals()["fn"+str(n)]  # string --> function

x = np.linspace(-10, 10, 40)
y = np.linspace(-10, 10, 40)
X, Y = np.meshgrid(x, y)    # numpy output required
Z = fn(X, Y)  

ax = plt.axes(projection='3d')
ax.plot_surface(X, Y, Z, 
    cmap="Blues_r", edgecolor="Black", lw=0.2)

plt.title('Surface Plot of fn'+str(n))
plt.xlabel('x')
plt.ylabel('y')
ax.set_zlabel('z')  # no plt.zlabel()
plt.show()