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

plt.style.use('seaborn')

def init_fn(x, mu, sigma):
    return np.exp(-(x-mu)**2 / (2*sigma**2)) / (np.sqrt(2*np.pi) * sigma)


L = 10     # length
T = 100     # time
c = 0.94
dx = 0.01
dt = 0.01
sigma2 = (c*dt/dx) ** 2

x = np.arange(0, L+dx, dx)
t = np.arange(0, T+dt, dt)
Nx = x.size
Nt = t.size

w = np.zeros((Nt, Nx))
w[:,0] = 1
w[:,-1] = 0
w[0,:] = init_fn(x, x.mean(), 0.05)
w[1,:] = w[0,:]

for itime in range(2, Nt):
    for ix in range(1, Nx-1):
        w[itime,ix] = 2*w[itime-1,ix] - w[itime-2,ix] + sigma2 * (w[itime-1,ix+1] - 2*w[itime-1,ix] + w[itime-1,ix-1])


# fig, axs = plt.subplots(nrows=4, ncols=3, figsize=(12, 10))
plt.ion()
idx = 0
for _ in range(len(w)):
    plt.cla()
    plt.plot(x, w[idx,:])
    # axs[i,j].plot(x, w[idx,:], 'k')
    # axs[i,j].set_title('time = {:.2f}'.format(t[idx]))
    plt.ylim([-10, 10])
    idx += 10
    plt.pause(0.01)
    plt.draw()
        
plt.tight_layout()
print("hutao tianxiadiyi")