import numpy as np
import PyOSGPUP.wavelet as wav
import matplotlib.pyplot as plt

t = np.arange(-0.4 / 2, (0.4 / 2) + 0.004, 0.004)
# frequency
fgR = np.array([10])  # 1 freq
fgO = np.array([5, 10, 15, 20])  # 4 fred
fgK = np.array([5, 10])  # 2 fred
gR = wav.getRicker(fgR, t)
gO = wav.getOrmsby(fgO, t)
gK = wav.getKlauder(fgK, t)

plt.figure()
plt.plot(t, gR, linewidth=2, color="blue")
plt.plot(t, gO, linewidth=2, color="red")
plt.plot(t, gK, linewidth=2, color="green")
plt.title('Wavelet')
plt.xlabel('time [ms]')
plt.ylabel('amplitude')
plt.legend(['Ricker', 'Ormsby', 'Kaluder'])
plt.show()

import PyOSGPUP.segypy as segypy

filename = 'shotgather.sgy'

# Set verbose level
segypy.verbose = 1
SH = segypy.getSegyHeader(filename)

# Read Segy File
[Data, SH, STH] = segypy.readSegy(filename)

# Plot Segy filwe
scale = 1e-9

# wiggle plot
segypy.wiggle(Data, SH, 1e-9)

# image plot
segypy.image(Data, SH, scale)

from PyOSGPUP.synthe_seismo import plotLogsInteract, plotTimeDepth, plotSeismogram
import matplotlib.pyplot as plt

# Synthetic Seismogram
d = [0., 50., 100.]  # Position of top of each layer (m)
v = [500., 1000., 1500.]  # Velocity of each layer (m/s)
rho = [2000., 2300., 2500.]  # Density of each layer (kg/m^3)
wavtyp = 'RICKER'  # Wavelet type
wavf = 50.  # Wavelet Frequency
usingT = False  # Use Transmission Coefficients?

plotLogsInteract(d[1], d[2], rho[0], rho[1], rho[2], v[0], v[1], v[2])
plt.show()

plotTimeDepth(d, v)
plt.show()

plotSeismogram(d, rho, v, 50., wavA=1., noise=0., usingT=True, wavtyp='RICKER')
plt.show()