%matplotlib inline
from SeisCL import SeisCL
import matplotlib.pyplot as plt
import numpy as np

"""
Explanation: Outputting a movie
In this brief example, we show how to create a movie of a seismic shot with SeisCL
End of explanation
"""


seis = SeisCL()

# Constants for the modeling
seis.ND = 2
N = 200
seis.N = np.array([N, 2*N])
seis.dt = dt = 0.25e-03
seis.dh = dh = 2
seis.NT = NT = 1000
seis.seisout = 1
seis.f0 = 20


# Source and receiver positions
sx = seis.N[1]//2 * dh
sy = 0
sz = seis.N[0]//2 * dh
gx = np.arange(N//4 * dh, (N - N//4)*dh, dh)
gy = gx * 0
gz = gx * 0 + N//4*dh
gsid = gz * 0
gid = np.arange(0, len(gz))
seis.src_pos_all = np.stack([[sx], [sy], [sz], [0], [100]], axis=0)
seis.rec_pos_all = np.stack([gx, gy, gz, gsid, gid, gx * 0, gx * 0, gx * 0], axis=0)
  
# We start with a simple model
vp_a = np.zeros(seis.N) + 3500
vs_a = np.zeros(seis.N) + 2000
rho_a = np.zeros(seis.N) + 2000

"""
Explanation: We first create a constant velocity model, with one source in the middle
End of explanation
"""


seis.movout = 20

seis.set_forward(seis.src_pos_all[3, :], {"vp": vp_a, "rho": rho_a, "vs": vs_a}, withgrad=False)
stdout = seis.execute()

"""
Explanation: To output a movie, we have to set the input 'movout' to a number greater than zero. For movout=10, the movie will contain every 10 time steps.
End of explanation
"""


movs = seis.read_movie()

"""
Explanation: SeisCL python wrapper contains a method to read the movie file.
End of explanation
"""


from matplotlib import animation
from IPython.display import HTML

toplot = movs[0][:,:,:,0]
fig = plt.figure(figsize=(6, 6))
im = plt.imshow(toplot[:,:,0], animated=True, vmin=np.min(toplot) / 10, vmax=np.max(toplot) / 10)

def init():
    im.set_array(toplot[:,:,0])
    return im,

def animate(t):
    im.set_array(toplot[:,:,t])
    return [im]
plt.close()

anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=movs[0].shape[2]-1, interval=100, blit=True, repeat=True)
HTML(anim.to_html5_video())

"""
Explanation: This last variable contains a list of movies for all the ouput variables given by seisout. In our case, seisout=1, so the outputs are vx and vz. We can visualize the movie with the following code.
End of explanation
"""


