from rsfproj 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 |
     sfsmooth rect1=3
     ''')

# Plot model
Result('model',
       '''
       grey allpos=y title=Model
       scalebar=y barreverse=y barlabel=Velocity barunit=km/s 
       ''')

# Constant density
Flow('density','model','math output=1')

# Wavelet shift
dt = 0.001
nt = 201
tau = (nt-1)*dt

# Source wavelet
Flow('wavelet',None,
     '''
     spike nsp=1 n1=2000 d1=%g k1=%d |
     ricker1 frequency=10 |
     transp
     ''' % (dt,nt))

# Source location
x = 1
z = 0.5

Flow('source',None,
     'spike n1=3 nsp=3 k1=1,2,3 mag=%g,%g,1' % (x,z))

# Modeling
Flow('data wave','wavelet model source density',
     '''
     awefd verb=y free=n snap=y jsnap=5
     nb=50 tz=0.006125 nbx=50 tx=0.006125 
     vel=${SOURCES[1]}
     sou=${SOURCES[2]}
     rec=${SOURCES[2]}
     den=${SOURCES[3]}
     wfl=${TARGETS[1]}         
     ''')

# Movie of wave snapshots
Plot('wave',
     '''
     window f1=50 f2=50 f3=40 n1=201 n2=301 |
     grey gainpanel=all title=Wave
     label1=Depth unit1=km label2=Lateral unit2=km
     ''',view=1)

# Favorite time moment
##########
time = 1 # !!!!!!!! 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=%g zshot=%g | add add=%g' % (x,z,tau))

# Movie of first-arrival wavefronts
fronts = []
for snap in range(140):
    front = 'front%d' % snap
    fronts.append(front)
    tsnap = tau+10*snap*dt
    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' % time)

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

End()