from SeisCL import SeisCL
import numpy as np
import matplotlib.pyplot as plt
import os

"""
Explanation: Simple example
SeisCL is of program that performs forward and adjoint seismic modeling for imaging and inversion. In this first example, we show how to use SeisCL to perform basic seismic modeling with the python wrapper. This simple example will solve the 2D isotropic elastic wave equation. 
The first thing to do is to load the SeisCL class and relevant modules:
End of explanation
"""


seis = SeisCL()

"""
Explanation: SeisCL is controlled by a single class in python that will group all relevant information to perform forward and adjoint modeling. The class creates all parameters with default values: they can be set at inialization or can be changed later. Note that all the available parameters are defined in the docstring of the SeisCL class.
End of explanation
"""


help(seis.__init__)

"""
Explanation: The class creates all parameters with default values: they can be set at inialization or can be changed later. Note that all the available parameters are defined in the docstring of the SeisCL class, which can be accessed by using the help function:
End of explanation
"""


seis.ND = 2                                 # Number of dimension
seis.N = np.array([250, 250])               # Grid size [NZ, NX, NY]
seis.dh = dh = 2                            # Grid spatial spacing
seis.dt = dt = 0.25e-03                     # Time step size
seis.NT = NT = 1000                         # Number of time steps

"""
Explanation: Simulation Grid
We begin setting the simulation grid, a 2D square domain of size $\small 500 \times 500 \; m$ with a cell size of $2m$ in all directions. We perform 1000 timesteps with a step size of with  $0.25ms$ for a simulation time of $0.25s$.
End of explanation
"""


vp = 3500
vs = 2000
rho = 2000

vp_a = np.zeros(seis.csts['N']) + vp
vs_a = np.zeros(seis.csts['N']) + vs
rho_a = np.zeros(seis.csts['N']) + rho

model_dict = {"vp": vp_a, "rho": rho_a, "vs": vs_a} 

"""
Explanation: Velocity Model
For the elastic wave equation, three material parameters are required : $V_p$, $V_s$ and $\rho$. For the sake of simplicity, we the parameters are homegeneous throughout the model, with $V_p = 3500 m/s$, $V_s = 2000 m/s$, $\rho = 2000 kg/m^3$. We store all parameters in a dictionnary, which will be used by the SeisCL class.
End of explanation
"""


seis.surface_acquisition_2d(ds=1000)        

"""
Explanation: Sources and receivers
Source and receivers are totally custumizable by SeisCL: each source can have different receivers, and multiple sources can be fired simultaneously. This flexibilty comes with a price: the definition of the source and receivers becomes verbose pretty easily. 
We provide some utility methods to ease the definition of source and receivers. Here, we use the method SeisCL.surface_acquisition_2d, which create a surface acquisition geometry, with receivers and sources at a defined interval of the grid steps along all the model.
End of explanation
"""


print(seis.src_pos_all)

"""
Explanation: All this function did is to define the attributes SeisCL.src_pos_all and SeisCL.rec_pos_all.
End of explanation
"""


print(seis.rec_pos_all[:,0])

"""
Explanation: We see here that we have one source, at x=36m, z= 36m (just outside the absorbing boundary), and that the source type is 100 (explosive source). The format of this array is [x, y, z, srcid, src_type], where all shots having the same src_id are fired simultaneously.
End of explanation
"""


seis.set_forward([0], model_dict, withgrad=False)
seis.execute()
datafd = seis.read_data()[0]                             

"""
Explanation: We see here that the first receiver is at x=40m and z=36m. The receiver is for the source id 0, and have the receiver id 1. The receiver id should be unique and start at 1 for the first receiver. The format of this array is [x, y, z, srcid, recid, src_type, none, none, none].
Simulation
we can now perform the simulation. The method SeisCL.set_forward is used to initialize the simulation, i.e. writing the hdf5 files containing simulation parameters and the model. The method  SeisCL.execute() launches SeisCL (the compiled program) and wait for the end of execution.
End of explanation
"""


os.listdir("./seiscl")

"""
Explanation: All hdf5 files are created in a temporary directory, called seiscl by default
End of explanation
"""


clip = 0.1
extent = [min(seis.rec_pos_all[0]), max(seis.rec_pos_all[0]), (datafd.shape[0]-1)*dt, 0]
vmax = np.max(datafd) * clip
vmin = -vmax
fig, ax = plt.subplots(1, 1, figsize=[4, 6])
ax.imshow(datafd, aspect='auto', vmax=vmax, vmin=vmin, extent = extent,
                interpolation='bilinear', cmap=plt.get_cmap('Greys'))
ax.set_title("FD solution to elastic wave equation in 2D \n", fontsize=16, fontweight='bold')
ax.set_xlabel("Receiver position (m)")
ax.set_ylabel("time (s)")
plt.show()

"""
Explanation: Each file contains specific information required for simulations:
* SeisCL_csts.mat contains all simulation parameters (dt, dh, ...)
* SeisCL_model.mat contains the material parameter arrays 
* SeisCL_dout.mat contains the simulated data.
Visualization
With matplotlib, we can finally analyze and visualize the result of our simulation.
End of explanation
"""


