#!-/usr/bin/python
# PDE Solver, Wave equation

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

"""
The plan to plot wave equation
We solved using an ansatz u = A(k)^n*e^(ijkdel(x))
And we found an A(k), which is dependent on time and space steps
so we calculate A(k) based on these and plug it into u
"""

global del_x, del_t, k, mu, steps
del_x = 0.05
k = 1
mu = .5
# Try calculating del_t from stability requirement calcs
# this might not work
#del_t = del_x/(mu*np.sqrt(np.cos(j*k*del_x)-1))
del_t = 0.05
t_steps = 100
x_steps = 100
#Beta = 2 + mu**2*(del_t/del_x)**2*(2*np.cos(j*k*del_x)-2)
#A = Beta/2 + np.sqrt(Beta**2 - 4)/2

u = np.zeros((t_steps, x_steps))
# u[n][j] = u[time][space]
# Set inintial conditions
u[0][0] = 0
u[0][1] = 0
u[1][0] = 1
u[1][1] = 1

# Step through space
# then step through time
for n in range(1, t_steps-2):
	n = n+1
	for j in range(1, x_steps -1):
		u[n+1][j] = 2*u[n][j] - u[n-1][j] + (mu*del_t/del_x)**2*(u[n][j+1] - 2*u[n][j] + u[n][j-1])
		print u[n][j]

## Animation below was pulled from either Sam, Will or a code example found on the internets

 #First set up the figure, the axis, and the plot element we want to animate
fig = plt.figure()
ax = fig.add_subplot(111,xlim=(0, x_steps*del_x), ylim=(-2, 2))
ax.axis("on")
point, = ax.plot([], [],'b')

print len(u[:])
print len(u[1][:])

# initialization function: plot the background of each frame
def init_ani():
    point.set_data([], [])
    return point,
 
# animation function.  This is called sequentially
def animate(i):
	x = np.linspace(0, x_steps*del_x, x_steps)
	y = u[i,:]
	point.set_data(x,y)
	return point,

# call the animator.  blit=True means only re-draw the parts that have changed.
anim = animation.FuncAnimation(fig, animate, init_func=init_ani,
                               frames=t_steps, interval=20, blit=True)
 
# this is how you save your animation to file:
anim.save('1d_wave_damped.mp4', writer='ffmpeg', fps=30)
 
plt.show()