from rsf.proj import *

# Make a velocity model with a hyperbolic reflector
Flow('model',None,
     '''
     math n1=301 o1=-1 d1=0.01 output="sqrt(1+x1*x1)" |
     unif2 n1=201 d1=0.01 v00=1,2 |
     put label1=Depth unit1=km label2=Lateral unit2=km 
     label=Velocity unit=km/s | 
     smooth rect1=3
     ''')

# Plot model
Result('model',
       '''
       grey allpos=y title=Model bias=1
       scalebar=y barreverse=y 
       ''')

# Source wavelet
Flow('wavelet',None,
     '''
     spike nsp=1 n1=2000 d1=0.001 k1=201 |
     ricker1 frequency=10
     ''')

# Extended model (for absorbing boundaries)
Flow('left','model',
     '''
     window n2=1 | spray axis=2 n=50 o=-1.5 d=0.01 |
     math output="input*exp(-4*(-1-x2)^2)"
     ''')
Flow('right','model',
     '''
     window n2=1 f2=300 | spray axis=2 n=50 o=3.01 d=0.01 |
     math output="input*exp(-4*(3-x2)^2)"
     ''')
Flow('emodel2','left model right',
     'cat axis=2 ${SOURCES[1:3]}')

Flow('top','emodel2',
     '''
     window n1=1 | spray axis=1 n=50 o=-0.5 d=0.01 |
     math output="input*exp(-4*x1^2)"
     ''')
Flow('bottom','emodel2',
     '''
     window n1=1 f1=200 | spray axis=1 n=50 o=2.01 d=0.01 |
     math output="input*exp(-4*(2-x1)^2)"
     ''')
Flow('emodel','top emodel2 bottom',
     'cat axis=1 ${SOURCES[1:3]}')

# Source location
Flow('source',None,
     '''
     spike k1=101 k2=251
     n2=401 o2=-1.5 d2=0.01 label1=Depth   unit1=km
     n1=301 o1=-0.5 d1=0.01 label2=Lateral unit2=km
     ''')

# Modeling
exe = Program('wave.c')
Flow('wave','source %s wavelet emodel' % exe[0],
     '''
     ./${SOURCES[1]} wav=${SOURCES[2]} v=${SOURCES[3]} 
     jt=5 ft=200
     ''')

# Movie of wave snapshots
Plot('wave',
     '''
     window f1=50 f2=50 n1=201 n2=301 | 
     grey gainpanel=all title=Wave
     ''',view=1)

# Your favorite time moment
###########################
time = 1.0 # !!! MODIFY ME
###########################

# Wavefield snapshot
Plot('snap','wave',
     '''
     window f1=50 f2=50 n1=201 n2=301 n3=1 min3=%g |
     grey title="Wave Snapshot at %g s"
     label1=Depth unit1=km label2=Lateral unit2=km
     ''' % (time,time))

# First-arrival traveltime
Flow('first','model',
     'eikonal yshot=1 zshot=0.5 | add add=0.2')

# Movie of first-arrival wavefronts
fronts = []
for snap in range(180):
    front = 'front%d' % snap
    fronts.append(front)
    tsnap = 0.2+snap*0.01
    Plot(front,'first',
         'contour nc=1 c0=%g title="%g s" ' % (tsnap,tsnap))
Plot('fronts',fronts,'Movie',view=1)

# First-arrival wavefront
Plot('front','first',
     '''
     contour nc=1 c0=%g wanttitle=n wantaxis=n
     plotcol=3 plotfat=5 
     ''' % time)

# Overlay wavefront and traveltime
Result('snap','snap front','Overlay')

End()