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

# k determines the size of the convection term (i.e.: movement speed)
k = 0.0

# c is (the square root of) the diffusivity
c = 0.5


def f(x,t):
    if t <= 0:
        result = 0*x
        return result
    else:
        result = np.exp(-((x+k*t)**2)/(4*t*(c**2)))
        result = result/(2*c*np.sqrt(np.pi*t))
        return result


# define the (discretised) region of interest
x = np.arange(-3, 3, 0.01)

# create an empty figure
fig = plt.figure()
# define the size of the figure
ax = plt.axes(xlim = (-3,3), ylim = (-0.2,5))
# plot an (empty) line - content comes later
line, = ax.plot([], [], lw=2)

# initialize the plot (background)
# the object that is updated in each frame has to be returned by this function
def init():
    line.set_data([], [])
    return line,

# define the frames of the animation
# again, the function has to return the object that change
def animate(i):
    # here 
    line.set_data(x,(f(x,0.001*i))) 
    return line,


# define the routine that creates the animation
# np.arange...basically defines the length of the animation; replace 1000 by some
#             larger number for even longer animations
# interval...time interval between different frames
# blit...if true, only the parts of the animation that have changed are redrawn
#        may speed up the animation
# the animation will be created online - here we define only its 'constructor'
ani = animation.FuncAnimation(fig, animate, np.arange(1, 1000), init_func=init,
    interval=50, blit=True)

# if you want to save the animations (as ogg with 30 frames per second), uncomment
# the following code - you need to have ffmpeg.exe installed, thus this code is
# not guaranteed to work on every computer
ani.save('anim.mp4',fps=30,writer="mencoder")


# run the animation
plt.show()