import numpy as np

# to plot the surface
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D



def phi(x,t):
    """
    compute some value for 
    for x and t which must be 1D array each
    """
    result = np.empty( [t.size,x.size] )
    L = x[-1]

    for it in range(t.size):
        result[it,:] = x*(L-x) * (1-it/t.size)

    return result



# setting some data
L    = 3.
Nx   = 51
Tmax = 3600.
Nt   = 361

# discretize bar of length L along x
list_x = np.linspace(0.,L,Nx)
# discretize time
list_t = np.linspace(0.,Tmax,Nt)

# compute phi for all values in list_x
phi_list = phi(list_x, list_t)

print(f"shape of phi : {phi_list.shape}")
print(f"shape of x   : {list_x.shape}")
print(f"shape of t   : {list_t.shape}")

X, T = np.meshgrid(list_x, list_t)
print(f"shape of X   : {X.shape}")
print(f"shape of T   : {T.shape}")


# plot the surface
fig = plt.figure()
ax = fig.gca(projection='3d')

ax.plot_surface(X,T, phi_list)
plt.show()