from rsf.proj import *

# set up main parameters
nt = 900; dt = 0.002      # time sampling
nz = 201; dz = 10; oz = 0 # spatial sampling in z
nx = 401; dx = 10; ox = 0 # spatial sampling in x
nb = 30; decay = 0.01     # ABC parameters
snap = 5                  # snapshot interval

# Velocity model construction
Flow('bvel1.asc',None,
     '''
     echo
     0    700
     500  900
     1000 1000
     1500 900
     2000 700
     2500 900
     3000 1000
     3500 900
     4000 700
     n1=2 n2=9 in=$TARGET
     data_format=ascii_float
     ''')
Flow('bvel2.asc',None,
     '''
     echo
     0 1500
     1000 1500
     2000 1500
     3000 1500
     4000 1500
     n1=2 n2=5 in=$TARGET
     data_format=ascii_float
     ''')

Flow('lay1','bvel1.asc',
     '''
     dd form=native |
     spline n1=401 d1=10 o1=0
     ''')
Flow('lay2','bvel2.asc',
     '''
     dd form=native |
     spline n1=401 d1=10 o1=0
     ''')
Flow('lays','lay1 lay2','cat axis=2 ${SOURCES[1:2]}')

Flow('vel1','lays',
     '''
     unif2 n1=%d d1=%g o1=%g v00=2000,2500,3000 |
     put label1=Depth unit1=m label2=Distance unit2=m
     label=Velocity unit=m/s
     ''' % (nz,dz,oz) )
Flow('leftv','vel1',
    '''
    window n2=1 |spray axis=2 n=200 d=10 o=-2000
    ''')
Flow('rightv','vel1',
    '''
    window n2=1 f2=400 |spray axis=2 n=200 d=10 o=4010
    ''')
Flow('vel','leftv vel1 rightv','cat axis=2 ${SOURCES[1:3]} | put o2=0')
Result('vel','grey title=Velocity color=j scalebar=y barreverse=y allpos=y')

# Get reflectivity from velocity
Flow('ref','vel',
    '''
    depth2time velocity=$SOURCE nt=1000 dt=0.002 |
    ai2refl |ricker1 frequency=10 |
    time2depth velocity=$SOURCE
    ''')
Result('ref','grey title=Reflectivity')

Flow('migvel','vel','smooth rect1=5 rect2=5 repeat=3 | math output="input/2" ')
Result('migvel','grey title="Migration Velocity" color=j scalebar=y barreverse=y allpos=y')

# Executable program
proj = Project()
prog = proj.Program('zofdrtm2.c')

# Forward exploding reflector modeling to get zero-offset data
Flow('dat wave','ref %s migvel' % prog[0],
     '''
     ./${SOURCES[1]} vel=${SOURCES[2]} verb=y nt=%d dt=%g mig=n nb=%d c=%g snap=%d snaps=${TARGETS[1]}
     ''' %(nt,dt,nb,decay,snap) )

Result('dat','grey title="Zero-offset Data" ')
Plot('wavem','wave','window j3=2 | grey gainpanel=all title=Forward')

# Exploding reflector migration to get image
Flow('img wave2','dat %s migvel' % prog[0],
     '''
     ./${SOURCES[1]} vel=${SOURCES[2]} verb=y mig=y nb=%d c=%g snap=%d snaps=${TARGETS[1]}
     ''' %(nb,decay,snap) )

Result('img','grey title="RTM Image" unit1=m')
Plot('wave2m','wave2','window j3=2 | grey gainpanel=all title=Backward')

End()