from rsf.proj import *
from math import pi

# Make a reflector model 
Flow('refl',None,
     '''
     math n1=10001 o1=-4 d1=0.001 output="sqrt(1+x1*x1)" 
     ''')
Flow('model','refl',
     '''
     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
     ''')
Result('model',
       '''
       window min2=-1 max2=2 j2=10 |
       grey allpos=y title=Model 
       scalebar=y barreverse=y 
       ''')

# Reflector dip
Flow('dip','refl','math output="x1/input" ')

# Kirchoff modeling
Flow('shot','refl dip',
     '''
     kirmod nt=1501 twod=y vel=1 freq=10  
     ns=1 s0=1 ds=0.01 nh=801 dh=0.01 h0=-4
     dip=${SOURCES[1]} verb=y |
     costaper nw2=200
     ''')
Plot('shot',
     '''
     window min2=-2 max2=2 min1=1 max1=4 |
     grey title="0 km Depth" 
     label1=Time unit1=s label2=Offset unit2=km
     labelsz=10 titlesz=15 grid=y
     ''')

# Double Fourier transform

Flow('kw','shot','fft1 | fft3 axis=2')

# Extrapolation filters

dx = 0.01
dz = 1

dx2 = 2*pi*dx
dz2 = 2*pi*dz
a = 0.5*dz2/(dx2*dx2)

# In the expressions below,
# x1 refers to omega and x2 refers to k

Flow('exact','kw',
     'math output="exp(I*sqrt(x1*x1-x2*x2)*%g)" ' % dz2)

Flow('approximate','kw',
     '''
     window f1=1 |
     math output="exp(I*x1*%g)* 
     (x1+I*(cos(x2*%g)-1)*%g)/
     (x1-I*(cos(x2*%g)-1)*%g)" |
     pad beg1=1 
     ''' % (dz2,dx2,a,dx2,a))

# Extrapolation

for case in ('exact','approximate'):
    Flow('phase-'+case,case,
         '''
         window n1=1 min1=5 min2=-2 max2=2 | 
         math output="log(input)" | sfimag
         ''')
    Flow('angle-'+case,'phase-'+case,
         '''
         math output="%g*asin(x1/5)" |
         cmplx $SOURCE
         ''' % (180/pi))

    Flow('shot-'+case,['kw',case],
         '''
         mul ${SOURCES[1]} | 
         fft3 axis=2 inv=y | fft1 inv=y
         ''')
    Plot('shot-'+case,
         '''
         window min2=-2 max2=2 min1=1 max1=4 |
         grey title="%g km Depth (%s)" 
         label1=Time unit1=s label2=Offset unit2=km
         labelsz=10 titlesz=15 grid=y
         ''' % (dz, case.capitalize()))

Result('shot','shot shot-exact shot-approximate',
       'SideBySideAniso')

Result('phase','angle-exact angle-approximate',
       '''
       cat axis=2 ${SOURCES[1]} | 
       graph title=Phase dash=0,1 
       label1="Angle From Vertical" unit1="\^o\_"
       ''')

End()