from numpy import exp,sin,linspace,array,eye,dot,empty_like,empty,flipud,pi
from matplotlib.pyplot import plot,figure,ion,ioff,draw,pause,imshow,imsave,show,title
from waterfall import waterfall

Tjiggle = 0.2
c = 1.
sigma = 1.0001 #.5 # 1. #1.000011
sigmasq = sigma**2
L = 1.
M = 1000
h = L/float(M)  # spatial grid spacing
print('pi/h =',pi/h)
#quit()
T = 0.7 
k = sigma*h/c  # time step
N = int(T/k)
print('Will take',N,'time steps of size k =', k)

m = M - 1

def l(t): return exp(-(t-Tjiggle)**2/2/(0.25*Tjiggle)**2)*(sin(100*t))#+0.3*sin(1500*t))  # jiggle left end for a little while  # add high-frequeny component to see numerical dispersion
def r(t): return 0.*t
def f(x): return 0.*x
def g(x): return 0.*x

x = linspace(0,L,M+1)
w = empty((N+1,M+1))
t = array([j*k for j in range(N+1)])

# write in boundary conditions

w[:,0] = l(t)
w[:,-1] = r(t)


# first step different
w[0,:] = f(x)  # initial value
w[1,1:-1] = w[0,1:-1] + k*g(x[1:-1]) + sigmasq/2*( w[0,:-2] - 2*w[0,1:-1] + w[0,2:] )  # account for initial velocity, g(x)

figure(0,figsize=(15,8))
ion()
plot([0,0],[-3,3],'w') # invisible line to set vertical scale
wave,=plot(x,w[0])
wave.set_ydata(w[1])

# main loop
for j in range(1,N):

	w[j+1,1:-1] = 2*w[j,1:-1] - w[j-1,1:-1] + sigmasq*( w[j,:-2] - 2*w[j,1:-1] + w[j,2:] )

	wave.set_ydata(w[j+1])
	draw()
	pause(0.0001) # needed for animation to show
	if abs(w[j+1,:]).max() > 3: break
ioff()

title('$\sigma='+str(sigma)+'$')
# x,t plane with u indicated by color
figure(1)
imshow(flipud(w[:j+1,:]),interpolation='nearest',cmap='rainbow')  # flip for time increasing upwards
from waterfall import waterfall
imsave( "wave.png",flipud(w[:j+1,:]),cmap='rainbow' )

# make a waterfall plot
figure(2)
jskip = 10
wsample = w[:j+1:jskip,:]
waterfall(x,range(wsample.shape[0]),wsample,slabthickness=3,zrange=[-4,6])
show()