from rsf.proj import *

# Generate a reflector model

layers = (
     ((0,2),(3.5,2),(4.5,2.5),(5,2.25),
       (5.5,2),(6.5,2.5),(10,2.5)),
     ((0,2.5),(10,3.5)),
     ((0,3.2),(3.5,3.2),(5,3.7),
       (6.5,4.2),(10,4.2)),
     ((0,4.5),(10,4.5))
)

nlays = len(layers)
for i in range(nlays):
     inp = 'inp%d' % i
     Flow(inp+'.asc',None,
          '''
          echo %s in=$TARGET
          data_format=ascii_float n1=2 n2=%d
          ''' % \
          (' '.join(map(lambda x: ' '.join(map(str,x)),
                           layers[i])),len(layers[i])))

dim1 = 'o1=0 d1=0.01 n1=1001'
Flow('lay1','inp0.asc',
     'dd form=native | spline %s fp=0,0' % dim1)
Flow('lay2',None  ,
     'math %s output="2.5+x1*0.1" '      % dim1)
Flow('lay3','inp2.asc',
     'dd form=native | spline %s fp=0,0' % dim1)
Flow('lay4',None  ,'math %s output=4.5'  % dim1)

Flow('lays','lay1 lay2 lay3 lay4',
     'cat axis=2 ${SOURCES[1:4]}')

graph = '''
graph min1=2.5 max1=7.5 min2=0 max2=5
yreverse=y wantaxis=n wanttitle=n screenratio=1
'''
Plot('lays0','lays',graph + ' plotfat=10 plotcol=0')
Plot('lays1','lays',graph + ' plotfat=2 plotcol=7')
Plot('lays2','lays',graph + ' plotfat=2')

# Velocity

Flow('vofz',None,
     '''
     math output="1.5+0.25*x1"
     d2=0.05 n2=201 d1=0.01 n1=501
     label1=Depth unit1=km
     label2=Distance unit2=km
     label=Velocity unit=km/s
     ''')
Plot('vofz',
     '''
     window min2=2.75 max2=7.25 |
     grey color=j allpos=y bias=1.5
     title=Model screenratio=1
     ''')

Result('model','vofz lays0 lays1','Overlay')

# Model data

Flow('dips','lays','deriv | scale dscale=100')
Flow('modl','lays dips',
     '''
     kirmod cmp=y dip=${SOURCES[1]} 
     nh=51  dh=0.1  h0=0
     ns=201 ds=0.05 s0=0
     freq=10 dt=0.004 nt=1501
     vel=1.5 gradz=0.25 verb=y |
     tpow tpow=1 |
     put d2=0.05 label3=Midpoint unit3=km 
     ''',split=[1,1001], reduce='add')

# Add random noise
Flow('data','modl','noise var=1e-6 seed=101811')

Result('data',
       '''
       byte |
       transp plane=23 |
       grey3 flat=n frame1=750 frame2=100 frame3=10 
       label1=Time unit1=s 
       label3=Half-Offset unit3=km 
       title=Data point1=0.8 point2=0.8
       ''')

# Initial constant-velocity migration
#####################################
Flow('mig','data',
     '''
     transp plane=23 |
     spray axis=3 n=1 d=0.1 o=0 |
     preconstkirch vel=1.5 |
     halfint inv=1 adj=1
     ''',split=[2,51],reduce='cat axis=4')

Result('mig',
       '''
       byte | window |
       grey3 flat=n frame1=750 frame2=100 frame3=10 
       label1=Time unit1=s 
       label3=Half-Offset unit3=km 
       title="Initial Migration" point1=0.8 point2=0.8
       ''')

# Velocity continuation
#######################

Flow('thk','mig',
     'window | transp plane=23 | cosft sign3=1')
Flow('velconk','thk',
     'fourvc nv=81 dv=0.01 v0=1.5 verb=y',
     split=[3,201])
Flow('velcon','velconk','cosft sign3=-1')

Plot('velcon',
     '''
     transp plane=23 memsize=1000 |
     window min2=2.5 max2=7.5 |
     grey title="Velocity Continuation"
     ''',view=1)

# Continue data squared
Flow('thk2','mig',
     '''
     mul $SOURCE |
     window | transp plane=23 | cosft sign3=1
     ''')
Flow('velconk2','thk2',
     'fourvc nv=81 dv=0.01 v0=1.5 verb=y',
     split=[3,201])
Flow('velcon2','velconk2','cosft sign3=-1')

# Compute semblance
Flow('semb','velcon velcon2',
     '''
     mul $SOURCE | divn den=${SOURCES[1]} rect1=25
     ''',split=[3,201])

Plot('semb',
     '''
     byte gainpanel=all allpos=y |
     transp plane=23 |
     grey3 flat=n frame1=750 frame2=0 frame3=48 
     label1=Time unit1=s color=j 
     label3=Velocity unit3=km/s movie=2 dframe=5
     title=Semblance point1=0.8 point2=0.8
     ''',view=1)

# Extracting images
###################
Flow('voft','vofz',
     'depth2time velocity=$SOURCE dt=0.004 nt=1501')
Flow('vrms','voft',
     '''
     add mode=p $SOURCE | causint | 
     math output="sqrt(input*0.004/(x1+0.004))" 
     ''')

# Using vrms is CHEATING
########################
Flow('slice','velcon vrms','slice pick=${SOURCES[1]}')

# Using vofz is CHEATING
########################
Flow('dmig','slice vofz',
     'time2depth velocity=${SOURCES[1]}')

Plot('dmig',
     '''
     window max1=5 min2=2.5 max2=7.5 | 
     grey title="Time -> Depth" screenratio=1
     label2=Distance label1=Depth unit1=km
     ''')

Result('dmig','Overlay')
Result('dmig2','dmig lays2','Overlay')

End()