##################################################################
# 
# Script used to save data generated by Madagascar into .npz files
#
# Run: python savedata.py kind
#
# where kind=_full, _nofs_full, _nosea_full
#
# M. Ravasi, 07/02/2022
#
##################################################################
import m8r
import sys
import numpy as np
import matplotlib.pyplot as plt


def read_shotgather(filepath, verb=False):
    if verb: print('Reading %s' % filepath)
    Rshot_rsf = m8r.File(filepath)
    Rshot = np.array(Rshot_rsf[:]).reshape((2351, 3, 180))
    return Rshot

datapath='./'

# Geometry
r_rsf = m8r.File(datapath+'r.rsf')
r = np.array(r_rsf[:]).reshape(r_rsf.int('n2'), r_rsf.int('n1')).T 
nr = r.shape[1] 
dr = r_rsf.int('d2')

# Sources
s_rsf = m8r.File(datapath+'s.rsf')
s = np.array(s_rsf[:]).reshape(s_rsf.int('n2'), s_rsf.int('n1')).T 
ns = s.shape[1] 
ds = s_rsf.int('d2')

# Velocity model
vel_rsf = m8r.File(datapath+'vp.rsf')
vel = np.array(vel_rsf[:]).reshape(vel_rsf.int('n2'), vel_rsf.int('n1')).T

z = np.arange(vel_rsf.int('n1'))*vel_rsf.float('d1') + vel_rsf.float('o1') 
x = np.arange(vel_rsf.int('n2'))*vel_rsf.float('d2') + vel_rsf.float('o2') 

# Wavelet
wav = np.array(m8r.File(datapath+'wav_ricker.rsf')[:][:1000:4])

# Time axis
ot, dt, nt =  0, 0.002, 2351
t = np.arange(nt)*dt + ot 

# Data
kind = sys.argv[1]
if kind == '_nosea_full':
	ns = nr
R = np.array([read_shotgather(datapath+'dat%s%d.rsf' % (kind,ishot), verb=True) for ishot in range(0, ns)])
nt = R.shape[1]
p = R[:,:, 0, :]
vz = R[:,:, 2, :]

# Data flat wav
Rflat = read_shotgather(datapath+'dat_noseaflat_full90.rsf', verb=True)
Rflatref = read_shotgather(datapath+'dat_noseaflat_ref90.rsf', verb=True)
pflatfull = Rflat[:, 0, :]
vzflatfull = Rflat[:, 2, :]
pflatref = Rflatref[:, 0, :]
vzflatref = Rflatref[:, 2, :]

pflat = pflatfull-pflatref
vzflat = vzflatfull-vzflatref

# Visualization
plt.figure(figsize=(10,5))
plt.imshow(vel, cmap='gray', extent = (x[0], x[-1], z[-1], z[0]), interpolation=None)
plt.scatter(s[0, ], s[1, ], marker='*', s=30, c='r', edgecolors='k')
plt.scatter(r[0, ],  r[1, ], marker='v', s=30, c='b', edgecolors='k')
plt.axis('tight')
plt.xlabel('x [m]'),plt.ylabel('y [m]'),plt.title('Model and Geometry')
plt.xlim(x[0], x[-1])

fig, axs = plt.subplots(1, 2, figsize=(9, 5))
fig.suptitle('Data')
axs[0].imshow(p[0], cmap='gray', vmin=-1e3, vmax=1e3, extent=(r[0,0], r[0,-1], t[-1], t[0]))
axs[0].set_title('p'), axs[0].set_xlabel(r'$x_R$'), axs[0].set_ylabel(r'$t$')
axs[0].axis('tight')
axs[1].imshow(vz[0], cmap='gray', vmin=-1e-3, vmax=1e-3, extent=(r[0,0], r[0,-1], t[-1], t[0]))
axs[1].set_title('vz'), axs[1].set_xlabel(r'$x_R$'), axs[1].set_ylabel(r'$t$')
axs[1].axis('tight')

fig, axs = plt.subplots(1, 3, figsize=(9, 5))
fig.suptitle('Data')
axs[0].imshow(pflatfull, cmap='gray', vmin=-1e5, vmax=1e5, extent=(r[0,0], r[0,-1], t[-1], t[0]))
axs[0].set_title('p'), axs[0].set_xlabel(r'$x_R$'), axs[0].set_ylabel(r'$t$')
axs[0].axis('tight')
axs[1].imshow(pflatref, cmap='gray', vmin=-1e5, vmax=1e5, extent=(r[0,0], r[0,-1], t[-1], t[0]))
axs[1].set_title('pref'), axs[1].set_xlabel(r'$x_R$'), axs[1].set_ylabel(r'$t$')
axs[1].axis('tight')
axs[2].imshow(pflat, cmap='gray', vmin=-1e5, vmax=1e5, extent=(r[0,0], r[0,-1], t[-1], t[0]))
axs[2].set_title('pref'), axs[1].set_xlabel(r'$x_R$'), axs[1].set_ylabel(r'$t$')
axs[2].axis('tight')

# Save
np.savez(datapath+'input%s_volvesynth' % kind, p=p, vz=vz, wav=wav, vel=vel, r=r, s=s, t=t, x=x, z=z)
np.savez(datapath+'input_noseaflat_volvesynth', pflat=pflat, vzflat=vzflat, r=r, s=s, t=t, x=x, z=z)

