import numpy as np
import matplotlib.pyplot as plt


def f(X, Y, t=0):
    return np.sin(X+t*np.pi)*np.sin(Y)

xmin, xmax = 0.0, 4*np.pi
ymin, ymax = 0.0, 5*np.pi
Nx, Ny = 50, 40
x, y = np.linspace(xmin, xmax, num=Nx), np.linspace(ymin, ymax, num=Ny)
X, Y = np.meshgrid(x, y, indexing='ij')
zstart = f(X, Y)

run_number = np.random.rand()

# Intialize figure stuff
fig = plt.figure()
axs = [
    plt.subplot(221, title='control flat'),
    plt.subplot(222, title='test flat'),
    plt.subplot(223, title='control gouraud'),
    plt.subplot(224, title='test gouraud'),
    ]

def drawcmesh(X, Y, z, ax, shading):
    ax.set_xlim(left=xmin, right=xmax)
    ax.set_ylim(bottom=ymin, top=ymax)
    cmesh = ax.pcolormesh(X, Y, z,
                          shading=shading, vmin=-1.5, vmax=1.5, cmap='hot')
    return cmesh



zstart = f(X, Y)
perm_mesh = [
    drawcmesh(X, Y, zstart, axs[1], shading='flat'),
    drawcmesh(X, Y, zstart, axs[3], shading='gouraud')
    ]
drawcmesh(X, Y, zstart, axs[0], shading='flat')
drawcmesh(X, Y, zstart, axs[2], shading='gouraud')

def update(t):
    z = f(X, Y, t)

    for i in range(4):
        ax = axs[i]
        if i<2:
            shading='flat'
        else:
            shading='gouraud'

        if i%2 == 0:
            # Control case: redraw by hand
            axs[i].clear()
            drawcmesh(X, Y, z, axs[i], shading=shading)
            axs[i].set_title('control {}'.format(shading))

        if i==1:
            # Flat shading

            # # The following would copy the syntax of plt.pcolormesh() but it
            # # results in an error
            # perm_mesh[0].set_array(z)

            # The following does not error out but the plot is later incorrect
            perm_mesh[0].set_array(z.ravel())

            # # The following works
            # perm_mesh[0].set_array(z[:-1, :-1].ravel())

        if i==3:
            # Gouraud shading

            # # The following would copy the syntax of plt.pcolormesh() but it
            # # results in an error
            # perm_mesh[1].set_array(z)

            # The following works
            perm_mesh[1].set_array(z.ravel())

            # # The following results in an error at draw time:

            
            # perm_mesh[1].set_array(z[:-1, :-1].ravel())

# # Uncomment this to check that the initial drawing for the test case is correct
# plt.show(block=True)

plt.show(block=False)
for t in np.linspace(0.0, 4.0, num=50):
    update(t)
    fig.canvas.draw()
    fig.canvas.flush_events()