import os
import sys

import pySW4 as sw4
import matplotlib.pyplot as plt
import numpy as np

"""
Explanation: rfile example - USGS Bay Area Velocity Model
The original input file can be found in the sw4-v1.1/examples/rfile directory, also on SW4's GitHub repository.
However to generate the output used in this notebook run the simulation again with
this input file: https://github.com/shaharkadmiel/pySW4/blob/master/tutorials/berkeley.sw4
End of explanation
"""


f = '../sw4-v1.1/examples/rfile/berkeley.rfile'
model = sw4.prep.rfileIO.read(f, block_number='all', verbose=True)

"""
Explanation: Read the USGS Bay Area Model
Using the rfileIO library to read in the model.
End of explanation
"""


model.plot_topography(cmap='terrain')

"""
Explanation: Plotting the topography of the model
End of explanation
"""


cs = model.get_cross_section(0, 12, 0, 12)
cs.plot()

"""
Explanation: Plotting a cross-section of the model
End of explanation
"""


z, properties = model.get_z_profile(4, 4)

labels = ('rho, kg/m^3', 'vp, m/s', 'vs, m/s', 'qp', 'qs')

fig, ax = plt.subplots()

for p, label in zip(properties.T, labels):
    p = np.ma.masked_equal(p, -999)
    ax.plot(p, z, label=label)
ax.invert_yaxis()
ax.xaxis.tick_top()
ax.set_xlabel('Material properties')
ax.xaxis.set_label_position('top')
ax.set_ylabel('Z, km')
ax.legend(loc=0)

"""
Explanation: Plotting material properties along a z profile
End of explanation
"""


f = '../sw4-v1.1/examples/rfile/results/berkeley.cycle=00000.z=0.topo.sw4img'
topo = sw4.read_image(f)
topo.plot(cmap='terrain')

"""
Explanation: Plotting the topography that SW4 actually sees
End of explanation
"""


