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

def surf_plot(dp,ny,nx,save_path,UorV='sqrt',non_dp=False):
    fig = plt.figure()
    ax = fig.add_subplot(projection='3d')

    # Make data
    # 生成网格数据
    X, Y = np.meshgrid(np.arange(nx), np.arange(ny))

    if non_dp:
        Z = Z_x = Z_y = dp.reshape(ny,nx)
    else:
        Z = dp.reshape(ny,nx,2)
        Z_x = Z[:,:,0].reshape(ny,nx)
        Z_y = Z[:,:,1].reshape(ny,nx)

    if UorV == 'u':
        ax.plot_surface(X, Y, Z_x, rstride = 1, cstride = 1, cmap='rainbow')
    elif UorV == 'v':
        ax.plot_surface(X, Y, Z_y, rstride = 1, cstride = 1, cmap='rainbow')
    else:
        ax.plot_surface(X, Y, Z_x*Z_x+Z_y*Z_y, rstride = 1, cstride = 1, cmap='rainbow')
        
    plt.savefig(save_path)
    plt.show()


if __name__=='__main__':
    print('nothing')
    # surf_plot(global_displace,nely+1,nelx+1,'u')
    # surf_plot(global_displace,nely+1,nelx+1,'v')
    # surf_plot(global_displace,nely+1,nelx+1,'sqrt')